from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

# Load model and tokenizer
model_name = "valhalla/t5-small-qg-hl"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Function to generate question
def generate_question(context, answer):
    if not context.strip() or not answer.strip():
        return "Please enter both context and answer."
    input_text = f"generate question: {context.replace(answer, '<hl> ' + answer + ' <hl>')}"
    inputs = tokenizer.encode(input_text, return_tensors="pt")
    outputs = model.generate(inputs, max_length=64, num_beams=4, early_stopping=True)
    question = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return question

# Gradio interface
iface = gr.Interface(
    fn=generate_question,
    inputs=[
        gr.Textbox(lines=5, label="Context or Paragraph"),
        gr.Textbox(lines=1, label="Answer (highlighted text)")
    ],
    outputs="text",
    title="🧠 AI Question Generator",
    description="Enter a paragraph and the answer you want highlighted. The app generates a relevant question."
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)