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

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

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

    with gr.Tab("POS"):
        gr.Markdown(" Description ")

        with gr.Row():
            prompt_POS = gr.Textbox(show_label=False, placeholder="Enter prompt")
            send_button_POS = gr.Button("Send", scale=0)

        gr.Markdown("Strategy 1 QA")
        with gr.Row():
            vicuna_chatbot1_POS = gr.Chatbot(label="vicuna-7b", live=True)
            llama_chatbot1_POS = gr.Chatbot(label="llama-7b", live=False)
            gpt_chatbot1_POS = gr.Chatbot(label="gpt-3.5", live=False)
        gr.Markdown("Strategy 2 Instruction")
        with gr.Row():
            vicuna_chatbot2_POS = gr.Chatbot(label="vicuna-7b", live=True)
            llama_chatbot2_POS = gr.Chatbot(label="llama-7b", live=False)
            gpt_chatbot2_POS = gr.Chatbot(label="gpt-3.5", live=False)
        gr.Markdown("Strategy 3 Structured Prompting")
        with gr.Row():
            vicuna_chatbot3_POS = gr.Chatbot(label="vicuna-7b", live=True)
            llama_chatbot3_POS = gr.Chatbot(label="llama-7b", live=False)
            gpt_chatbot3_POS = gr.Chatbot(label="gpt-3.5", live=False)
        
    with gr.Tab("Chunk"):
        gr.Markdown(" Description 2 ")

        with gr.Row():
            prompt_chunk = gr.Textbox(show_label=False, placeholder="Enter prompt")
            send_button_Chunk = gr.Button("Send", scale=0)

        gr.Markdown("Strategy 1 QA")
        with gr.Row():
            vicuna_chatbot1_chunk = gr.Chatbot(label="vicuna-7b", live=True)
            llama_chatbot1_chunk = gr.Chatbot(label="llama-7b", live=False)
            gpt_chatbot1_chunk = gr.Chatbot(label="gpt-3.5", live=False)
        gr.Markdown("Strategy 2 Instruction")
        with gr.Row():
            vicuna_chatbot2_chunk = gr.Chatbot(label="vicuna-7b", live=True)
            llama_chatbot2_chunk = gr.Chatbot(label="llama-7b", live=False)
            gpt_chatbot2_chunk = gr.Chatbot(label="gpt-3.5", live=False)
        gr.Markdown("Strategy 3 Structured Prompting")
        with gr.Row():
            vicuna_chatbot3_chunk = gr.Chatbot(label="vicuna-7b", live=True)
            llama_chatbot3_chunk = gr.Chatbot(label="llama-7b", live=False)
            gpt_chatbot3_chunk = gr.Chatbot(label="gpt-3.5", live=False)
        
        clear = gr.ClearButton([prompt_chunk, vicuna_chatbot1_chunk])

    # Define the function for generating responses
    def generate_response(prompt):
        input_ids = tokenizer.encode(prompt, return_tensors="pt")
        output_ids = model.generate(input_ids, max_length=500, pad_token_id=tokenizer.eos_token_id)
        response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        return response

    # Define the Gradio interface
    def chatbot_interface_POS(input_dict):
        prompt_POS = input_dict["prompt_POS"]
        vicuna_response_POS = generate_response(prompt_POS)
        # Add responses from other chatbots if needed
        return {"Vicuna-7B": vicuna_response_POS}

    def chatbot_interface_Chunk(input_dict):
        prompt_chunk = input_dict["prompt_chunk"]
        vicuna_response_chunk = generate_response(prompt_chunk)
        # Add responses from other chatbots if needed
        return {"Vicuna-7B": vicuna_response_chunk}

    # Connect the interfaces to the functions
    send_button_POS.click(chatbot_interface_POS, inputs={"prompt_POS": prompt_POS, "vicuna_chatbot1_POS": vicuna_chatbot1_POS})
    send_button_Chunk.click(chatbot_interface_Chunk, inputs={"prompt_chunk": prompt_chunk, "vicuna_chatbot1_chunk": vicuna_chatbot1_chunk})



demo.launch()