import gradio as gr
from transformers import pipeline
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "finetuned_phi2"
model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

num_new_tokens = 200  # change to the number of new tokens you want to generate

DESCRIPTION = """\
# Microsoft Phi2 Chatbot
\n The model is hosted on a CPU and inference takes a long time. Please feel free to duplicate the space and use it on a GPU"""


def generate(question, context, max_new_tokens = 200, temperature = 0.6):
  
    system_message = "You are a question answering chatbot. Provide a clear and detailed explanation"
    prompt = f"[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n {question} [/INST]" # replace the command here with something relevant to your task

    # Count the number of tokens in the prompt
    num_prompt_tokens = len(tokenizer(prompt)['input_ids'])
    # Calculate the maximum length for the generation
    max_length = num_prompt_tokens + max_new_tokens

    gen = pipeline('text-generation', model=model, tokenizer=tokenizer, max_length=max_length, temperature=temperature)
    result = gen(prompt)
    return (result[0]['generated_text'].replace(prompt, ''))


bbchatbot = gr.Chatbot(
    avatar_images=["logo/user_logo.png", "logo/bot_logo.png"], bubble_full_width=False, show_label=False, show_copy_button=True, likeable=True,)

examples = [["What are transformers?"], ["What are LLMs"], ["What is machine learning?"], ["How to write a good resume?"]]

additional_inputs =  additional_inputs=[gr.Slider(label="Max new tokens",minimum=100,maximum=500,step=10,value=num_new_tokens),
                                        gr.Slider(label="Temperature",minimum=0.1,maximum=4.0,step=0.1,value=0.6)]

chat_interface  = gr.ChatInterface(fn=generate,
                        additional_inputs=additional_inputs,
                        chatbot=bbchatbot,
                        title="",
                        examples=examples
                       )

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(value="Duplicate for private use", elem_id="duplicate-button")
    chat_interface.render()


demo.queue().launch(show_api=False)