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)

## Task 1
# msg = template_all.format(text)
template_all = '''Output the <Noun, Verb, Adjective, Adverb, Preposition/Subord, Coordinating Conjunction, Cardinal Number, Determiner, Noun Phrase, Verb Phrase, Adjective Phrase, Adverb Phrase, Preposition Phrase, Conjunction Phrase, Coordinate Phrase, Quantitave Phrase, Complex Nominal, Clause, Dependent Clause, Fragment Clause, T-unit, Complex T-unit, Fragment T-unit> in the following sentence without additional text in json format: "{}"'''
# msg = template_single.format(ents_prompt[eid], text)
template_single = '''Output any <{}> in the following sentence one per line without additional text: "{}"'''

## Task 2
prompt2_pos = '''POS tag the following sentence using Universal POS tag set without generating additional text: {}'''
prompt2_chunk = '''Do sentence chunking for the following sentence as in CoNLL 2000 shared task without generating addtional text: {}'''

## Task 3

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

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

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

        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])
    
    with gr.Tab("Chunk"):
        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

    prompt_POS.submit(respond, [template_all.format(prompt_POS), vicuna_S1_chatbot_POS], [template_all.format(prompt_POS), vicuna_S1_chatbot_POS])
    prompt_POS.submit(respond, [prompt2_pos.format(prompt_POS), vicuna_S2_chatbot_POS], [prompt2_pos.format(prompt_POS), vicuna_S2_chatbot_POS])
    prompt_POS.submit(respond, [prompt_POS, vicuna_S3_chatbot_POS], [prompt_POS, vicuna_S3_chatbot_POS])

    prompt_CHUNK.submit(respond, [template_all.format(prompt_CHUNK), vicuna_S1_chatbot_CHUNK], [template_all.format(prompt_CHUNK), vicuna_S1_chatbot_CHUNK])
    prompt_CHUNK.submit(respond, [prompt2_chunk.format(prompt_CHUNK), vicuna_S2_chatbot_CHUNK], [prompt2_chunk.format(prompt_CHUNK), vicuna_S2_chatbot_CHUNK])
    prompt_CHUNK.submit(respond, [prompt_CHUNK, vicuna_S3_chatbot_CHUNK], [prompt_CHUNK, vicuna_S3_chatbot_CHUNK])

demo.launch()