import gradio as gr
from transformers import pipeline

# Initialize the GPT-2 pipeline
pipe = pipeline("text-generation", model="gpt2")

theme = gr.themes.Soft()

# Function that generates text based on instruction-based prompting
def generate_text(input_instruction):
    # Use the input instruction to generate text
    generated_text = pipe(input_instruction, max_length=500)[0]['generated_text']
    return generated_text

# Define example instructions for testing
instruction_examples = [
    ("Describe the origin of the universe"),
    ("Explain the concept of artificial intelligence"),
    ("Describe the most common types of cancer"),
]

# Function that echoes the input text
#def echo_text(input_text):
#    return input_text

with gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(placeholder="Enter text here..."),
    outputs=gr.Textbox(),
    examples=instruction_examples, 
    live=False,
    title="LLM Evaluator with Linguistic Scrutiny",
    theme=theme
) as iface:
    blocks = gr.Blocks()
    
    with gr.Row():
        vicuna_model_selector = gr.Dropdown(["7b", "13b", "33b"], label="Vicuna Model", placeholder="Select model size")
        llama_model_selector = gr.Dropdown(["7B", "13B", "30B", "65B"], label="LLaMa Model", placeholder="Select model size")
        chatgpt_api_key = gr.Textbox(label="ChatGPT API Key", type="password", placeholder="Enter your API key")

    # Strategy 1 - QA-Based Prompting
    with gr.Accordion("Strategy 1 - QA-Based Prompting", style="font-weight: bold; font-size: 16px;"):
        with gr.Row():
            chatgpt_btn = gr.Button("ChatGPT")
            llama_btn = gr.Button("LLaMA")
            vicuna_btn = gr.Button("Vicuna")
            alpaca_btn = gr.Button("Alpaca")
            flant5_btn = gr.Button("Flan-T5")

    # Strategy 2 - Instruction-Based Prompting
    with gr.Accordion("Strategy 2 - Instruction-Based Prompting", style="font-weight: bold; font-size: 16px;"):
        with gr.Row():
            chatgpt_btn = gr.Button("ChatGPT")
            llama_btn = gr.Button("LLaMA")
            vicuna_btn = gr.Button("Vicuna")
            alpaca_btn = gr.Button("Alpaca")
            flant5_btn = gr.Button("Flan-T5")

    # Strategy 3 - Structured Prompting
    with gr.Accordion("Strategy 3 - Structured Prompting", style="font-weight: bold; font-size: 16px;"):
        with gr.Row():
            chatgpt_btn = gr.Button("ChatGPT")
            llama_btn = gr.Button("LLaMA")
            vicuna_btn = gr.Button("Vicuna")
            alpaca_btn = gr.Button("Alpaca")
            flant5_btn = gr.Button("Flan-T5")
    
    iface.launch()