import gradio as gr
from transformers import pipeline

# Initialize a dictionary of models for each strategy
models = {
    "GPT-2": "gpt2",
    "ChatGPT": "EleutherAI/gpt-neo-1.3",
    "LLaMa": "model_for_LLaMa",
    "Vicuna": "model_for_Vicuna",
    "Alpaca": "model_for_Alpaca",
    "Flan-T5": "model_for_Flan-T5",
}

# Define a function to generate text based on the selected model
def generate_text(input_instruction, selected_model):
    if selected_model in models:
        model_name = models[selected_model]
        pipe = pipeline("text-generation", model=model_name)
        generated_text = pipe(input_instruction, max_length=100, do_sample=True)[0]['generated_text']
        return generated_text
    else:
        return "Please select a model for this strategy."

# Define example instructions for testing
instruction_examples = [
    ("Write a short story about a cat."),
    ("Explain the concept of artificial intelligence."),
    ("Compose a poem about nature."),
]

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(placeholder="Enter instruction here..."),
    outputs=gr.Textbox(),
    examples=instruction_examples,
    live=True,
    title="Text Generation with Dynamic Model Selection",
)

# Additional input section 1 - User input
with gr.Row():
    user_input = gr.Textbox(placeholder="Enter your input...")

# Additional input section 2 - Strategy 1
with gr.Row():
    strategy1_selector = gr.Dropdown(list(models.keys()), label="Strategy 1 - QA-Based Prompting")

# Additional input section 3 - Strategy 2
with gr.Row():
    strategy2_selector = gr.Dropdown(list(models.keys()), label="Strategy 2 - Instruction-Based Prompting")

# Additional input section 4 - Strategy 3
with gr.Row():
    strategy3_selector = gr.Dropdown(list(models.keys()), label="Strategy 3 - Structured Prompting")

# Create a callback function for dynamic model selection
def update_model_and_generate_text(input_instruction, selected_model):
    if selected_model:
        selected_model_name = strategy1_selector.value if selected_model == "Strategy 1 - QA-Based Prompting" else (
            strategy2_selector.value if selected_model == "Strategy 2 - Instruction-Based Prompting" else
            strategy3_selector.value)
        iface.set_function(generate_text, inputs=[user_input, selected_model], outputs="outputs")
        return "Selected model: " + selected_model_name

# Add a submit button to trigger dynamic model selection
submit_button = gr.Button("Submit")

# Set the function for the Gradio interface to the update_model_and_generate_text function
iface.fn = update_model_and_generate_text

# Add the submit button to the interface
iface.add(submit_button)

iface.launch()