import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "anirudh-sub/debate_model_v2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate_text(prompt):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(input_ids, max_length=100, num_return_sequences=1)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text

iface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    live=True,
    title="Debate Model",
    description="This model generates text based on the input prompt using Llama.",
)

iface.launch()