import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
model_name = "Tahsin-Mayeesha/squad-bn-mt5-base2"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)  


import gradio as gr
def generate__questions(context,answer):
  text='answer: '+answer + ' context: ' + context
  text_encoding = tokenizer.encode_plus(
      text,return_tensors="pt"
  )
  model.eval()
  generated_ids =  model.generate(
    input_ids=text_encoding['input_ids'],
    attention_mask=text_encoding['attention_mask'],
    max_length=64,
    num_beams=5,
    num_return_sequences=1
  )
  
  return tokenizer.decode(generated_ids[0],skip_special_tokens=True,clean_up_tokenization_spaces=True).replace('question: ',' ')

demo = gr.Interface(fn=generate__questions, inputs=[gr.Textbox(label='Context'),
                                                    gr.Textbox(label='Answer')] ,
                                                    outputs=gr.Textbox(label='Question'),
                                                    title="Bangla Question Generation",
                                                    description="Get the Question from given Context and an Answer")
demo.launch()