import gradio as gr
from transformers import AutoTokenizer, T5ForConditionalGeneration

model_name = "allenai/t5-small-squad2-question-generation"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

def generate_questions(input_string, max_length=80, temperature=1.0, num_return_sequences=2,
                        num_beams=4, top_k=90, top_p=0.9):
    input_ids = tokenizer.encode(input_string, return_tensors="pt")
    res = model.generate(input_ids, max_length=max_length, num_return_sequences=num_return_sequences,
                        num_beams=num_beams, temperature=temperature, top_k=top_k, top_p=top_p)
    output = tokenizer.batch_decode(res, skip_special_tokens=True)
    return output

input_text = gr.inputs.Textbox(label="Enter some text:", default="Nicejob has increased our revenue 80% since signing up")
max_length = gr.inputs.Slider(10, 150, 80, label="Max Length")
temperature = gr.inputs.Slider(0.0, 1.0, 1.0, step=0.05, label="Temperature")
num_return_sequences = gr.inputs.Slider(1, 10, 2, label="Num Return Sequences")
num_beams = gr.inputs.Slider(1, 10, 4, label="Num Beams")
top_k = gr.inputs.Slider(0, 100, 90, label="Top-k")
top_p = gr.inputs.Slider(0.0, 1.0, 0.9, step=0.05, label="Top-p")
output_text = gr.outputs.Textbox(label="Generated questions:")

iface = gr.Interface(
    generate_questions,
    inputs=[input_text, max_length, temperature, num_return_sequences, num_beams, top_k, top_p],
    outputs=output_text,
    title="Question Generation",
    description="Generate questions from text using the T5-SQuAD2 model.",
)

iface.launch()