import gradio as gr
import transformers
import torch
from peft import PeftModel
import os

HF_TOKEN = os.environ.get("HF_TOKEN")

model_id = "JerniganLab/interviews-and-qa"
base_model = "meta-llama/Meta-Llama-3-8B-Instruct"

llama_model = transformers.AutoModelForCausalLM.from_pretrained(base_model)


pipeline = transformers.pipeline(
    "text-generation",
    model=llama_model,
    tokenizer=base_model,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device="cuda",
)

pipeline.model = PeftModel.from_pretrained(llama_model, model_id)

# def chat_function(message, history, system_prompt, max_new_tokens, temperature):
#     messages = [{"role":"system","content":system_prompt},
#                 {"role":"user", "content":message}]
#     prompt = pipeline.tokenizer.apply_chat_template(
#         messages,
#         tokenize=False,
#         add_generation_prompt=True,)
#     terminators = [
#         pipeline.tokenizer.eos_token_id,
#         pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")]
#     outputs = pipeline(
#         prompt,
#         max_new_tokens = max_new_tokens,
#         eos_token_id = terminators,
#         do_sample = True,
#         temperature = temperature + 0.1,
#         top_p = 0.9,)
#     return outputs[0]["generated_text"][len(prompt):]

def chat_function(message, history, max_new_tokens, temperature):
    SYSTEM_PROPMT = "I want you to embody a 30-year-old Southern Black woman graduate student who is kind, empathetic, direct, unapologetically Black, and who communicates predominantly in African American Vernacular English. I want you to act as a companion for graduate students who are enrolled in primarily white universities. As their companion, I want you to employ principles of cognitive behavioral therapy, the rhetoric of Black American digital spaces, and Black American humor in your responses to the challenges that students encounter with peers, faculty, or staff. I want you to engage in role-play with them, providing them a safe place to develop potential responses to microaggressions. I want you to help them seek resolutions for their problems."
    messages = [{"role":"system","content":SYSTEM_PROPMT},
                {"role":"user", "content":message}]
    prompt = pipeline.tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,)
    terminators = [
        pipeline.tokenizer.eos_token_id,
        pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")]
    outputs = pipeline(
        prompt,
        max_new_tokens = max_new_tokens,
        eos_token_id = terminators,
        do_sample = True,
        temperature = temperature + 0.1,
        top_p = 0.9,)
    return outputs[0]["generated_text"][len(prompt):]

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
# demo = gr.ChatInterface(
#     chat_function,
#     textbox=gr.Textbox(placeholder="Enter message here", container=False, scale = 7),
#     chatbot=gr.Chatbot(height=400),
#     additional_inputs=[
#         gr.Textbox("You are helpful AI", label="System Prompt"),
#         gr.Slider(100,4000, label="Max New Tokens"),
#         gr.Slider(0,1, label="Temperature")
#     ]
#     )

demo = gr.ChatInterface(
    chat_function,
    textbox=gr.Textbox(placeholder="Enter message here", container=False, scale = 7),
    chatbot=gr.Chatbot(height=400),
    additional_inputs=[
        gr.Slider(100,4000, label="Max New Tokens"),
        gr.Slider(0,1, label="Temperature")
    ]
    )



if __name__ == "__main__":
    demo.launch()