import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

peft_model_id = f"mdacampora/tax-convos-demo2"
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(conversation):
    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()