File size: 2,057 Bytes
0af5155
 
 
 
9314a2e
0af5155
 
 
 
 
 
 
 
 
 
 
 
61da192
 
 
 
 
14d5482
61da192
 
14d5482
b2eedf8
 
 
f375845
61da192
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0af5155
 
 
 
 
 
9f35321
 
 
0af5155
9f35321
 
 
 
 
 
 
 
 
 
 
 
 
 
acb3080
9f35321
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

peft_model_id = f"mdacampora/tax-convos-demo"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    return_dict=True,
    load_in_8bit=True,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id)

def make_inference(problem, answer):
    batch = tokenizer(
        problem,
        return_tensors="pt",
    )

    with torch.cuda.amp.autocast():
        output_tokens = model.generate(**batch, max_new_tokens=50)





# def make_inference(conversation, response):
#     conversation_history = conversation
#     response = ""
#     while True:
#         batch = tokenizer(
#             f"### Problem:\n{conversation_history}\n{response}",
#             return_tensors="pt",
#         )
#         with torch.cuda.amp.autocast():
#             output_tokens = model.generate(**batch, max_new_tokens=50)
#         new_response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
#         if new_response.strip() == "":
#             break
#         response = f"\n{new_response}"
#         conversation_history += response
#     return conversation_history


if __name__ == "__main__":
    # make a gradio interface
    import gradio as gr

    # gr.Interface(
    #     make_inference,
    #     [
            
    #         gr.inputs.Textbox(lines=1, label="Problem"),
    #     ],
    #     gr.outputs.Textbox( label="Transcript"),
    #     title="tax-convos-demo",
    #     description="trying to create a crude chat bot for tax services.",
    # ).launch()

gr.Interface(
    make_inference,
    [
        gr.inputs.Textbox(lines=5, label="Conversation"),
    ],
    gr.outputs.Textbox(label="Updated Conversation"),
    title="tax-convos-demo",
    description="Ask any tax-related questions you may have.",
).launch()