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()