import gradio as gr
from transformers import pipeline

# Load the Meta Llma 3.1 Instruct model for the initial argument
model1 = gr.load("models/microsoft/GRIN-MoE")

# Load a different model for the counter-argument
# We'll use GRIN as an example, but you can replace this with another suitable model
model2 = gr.load("models/microsoft/GRIN-MoE")

def generate_initial_argument(query):
    prompt = f"Provide a logical explanation for the following topic: {query}"
    response = model1(prompt)
    return response

def generate_counter_argument(query, initial_argument):
    prompt = f"Given the topic '{query}' and the argument '{initial_argument}', provide a well-reasoned counter-argument:"
    response = model2(prompt, max_length=200, num_return_sequences=1, temperature=0.7)[0]['generated_text']
    # Extract the counter-argument from the generated text
    counter_argument = response.split(prompt)[-1].strip()
    return counter_argument

def debate(query):
    initial_argument = generate_initial_argument(query)
    counter_argument = generate_counter_argument(query, initial_argument)
    return initial_argument, counter_argument

# Define the Gradio interface
iface = gr.Interface(
    fn=debate,
    inputs=gr.Textbox(lines=2, placeholder="Enter your question or topic for debate here..."),
    outputs=[
        gr.Textbox(label="Initial Argument (Meta Llma 3.1 Instruct)"),
        gr.Textbox(label="Counter-Argument (GRIN-MoE)")
    ],
    title="Two-Model Debate System",
    description="Enter a question or topic. Meta Llma 3.1 Instruct will provide an initial argument, and GRIN-MoE will generate a counter-argument.",
    examples=[
        ["What are the long-term implications of artificial intelligence on employment?"],
        ["Should governments prioritize space exploration or addressing climate change?"],
        ["Is genetic engineering in humans ethical for disease prevention?"]
    ]
)

# Launch the interface
iface.launch()