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

# Load the Vicuna 7B v1.3 LMSys model and tokenizer
model_name = "lmsys/vicuna-7b-v1.3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

template_single = '''Output any <{}> in the following sentence one per line: "{}"'''

linguistic_entities = [
    "Noun",
    "Determiner",
    "Noun phrase",
    "Verb phrase",
    "Dependent Clause",
    "T-units"
]

with gr.Blocks() as demo:
    gr.Markdown("# LLM Evaluator With Linguistic Scrutiny")

    
    gr.Markdown(" Description ")

    # Dropdown for linguistic entities
    entity_dropdown = gr.Dropdown(linguistic_entities, label="Select Linguistic Entity")
    prompt_POS = gr.Textbox(show_label=False, placeholder="Write a prompt and press enter")
    submit_btn = gr.Button(label="Submit")


    gr.Markdown("Strategy 1 QA-Based Prompting")
    with gr.Row():
        vicuna_S1_chatbot_POS = gr.Chatbot(label="vicuna-7b")
        llama_S1_chatbot_POS = gr.Chatbot(label="llama-7b")
        gpt_S1_chatbot_POS = gr.Chatbot(label="gpt-3.5")
    clear = gr.ClearButton([prompt_POS, vicuna_S1_chatbot_POS])
    gr.Markdown("Strategy 2 Instruction-Based Prompting")
    with gr.Row():
        vicuna_S2_chatbot_POS = gr.Chatbot(label="vicuna-7b")
        llama_S2_chatbot_POS = gr.Chatbot(label="llama-7b")
        gpt_S2_chatbot_POS = gr.Chatbot(label="gpt-3.5")
    clear = gr.ClearButton([prompt_POS, vicuna_S2_chatbot_POS])
    gr.Markdown("Strategy 3 Structured Prompting")
    with gr.Row():
        vicuna_S3_chatbot_POS = gr.Chatbot(label="vicuna-7b")
        llama_S3_chatbot_POS = gr.Chatbot(label="llama-7b")
        gpt_S3_chatbot_POS = gr.Chatbot(label="gpt-3.5")
    clear = gr.ClearButton([prompt_POS, vicuna_S3_chatbot_POS])
    
    
    # gr.Markdown(" Description ")

    # prompt_CHUNK = gr.Textbox(show_label=False, placeholder="Write a prompt and press enter")

    # gr.Markdown("Strategy 1 QA")
    # with gr.Row():
    #     vicuna_S1_chatbot_CHUNK = gr.Chatbot(label="vicuna-7b")
    #     llama_S1_chatbot_CHUNK = gr.Chatbot(label="llama-7b")
    #     gpt_S1_chatbot_CHUNK = gr.Chatbot(label="gpt-3.5")
    # clear = gr.ClearButton([prompt_CHUNK, vicuna_S1_chatbot_CHUNK])
    # gr.Markdown("Strategy 2 Instruction")
    # with gr.Row():
    #     vicuna_S2_chatbot_CHUNK = gr.Chatbot(label="vicuna-7b")
    #     llama_S2_chatbot_CHUNK = gr.Chatbot(label="llama-7b")
    #     gpt_S2_chatbot_CHUNK = gr.Chatbot(label="gpt-3.5")
    # clear = gr.ClearButton([prompt_CHUNK, vicuna_S2_chatbot_CHUNK])
    # gr.Markdown("Strategy 3 Structured Prompting")
    # with gr.Row():
    #     vicuna_S3_chatbot_CHUNK = gr.Chatbot(label="vicuna-7b")
    #     llama_S3_chatbot_CHUNK = gr.Chatbot(label="llama-7b")
    #     gpt_S3_chatbot_CHUNK = gr.Chatbot(label="gpt-3.5")
    # clear = gr.ClearButton([prompt_CHUNK, vicuna_S3_chatbot_CHUNK])

    # def respond(message, chat_history):
    #     input_ids = tokenizer.encode(message, return_tensors="pt")
    #     output_ids = model.generate(input_ids, max_length=50, num_beams=5, no_repeat_ngram_size=2)
    #     bot_message = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        
    #     chat_history.append((message, bot_message))
    #     time.sleep(2)
    #     return "", chat_history
    
    def respond_entities(entity, message, chat_history):
        prompt = template_single.format(entity, message)
        input_ids = tokenizer.encode(prompt, return_tensors="pt")
        output_ids = model.generate(input_ids, max_length=50, num_beams=5, no_repeat_ngram_size=2)
        bot_message = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        
        chat_history.append((message, bot_message))
        time.sleep(2)
        return entity, message, bot_message

    submit_btn.click(respond_entities, [entity_dropdown, prompt_POS, vicuna_S1_chatbot_POS], [entity_dropdown, prompt_POS, vicuna_S1_chatbot_POS])
    submit_btn.click(respond_entities, [entity_dropdown, prompt_POS, vicuna_S2_chatbot_POS], [entity_dropdown, prompt_POS, vicuna_S2_chatbot_POS])
    submit_btn.click(respond_entities, [entity_dropdown, prompt_POS, vicuna_S3_chatbot_POS], [entity_dropdown, prompt_POS, vicuna_S3_chatbot_POS])

    # prompt_CHUNK.submit(respond_entities, [prompt_CHUNK, vicuna_S1_chatbot_CHUNK], [prompt_CHUNK, vicuna_S1_chatbot_CHUNK])
    # prompt_CHUNK.submit(respond_entities, [prompt_CHUNK, vicuna_S2_chatbot_CHUNK], [prompt_CHUNK, vicuna_S2_chatbot_CHUNK])
    # prompt_CHUNK.submit(respond_entities, [prompt_CHUNK, vicuna_S3_chatbot_CHUNK], [prompt_CHUNK, vicuna_S3_chatbot_CHUNK])

demo.launch()