import transformers
import sentencepiece
import torch
import numpy as np

from transformers import T5ForConditionalGeneration,T5Tokenizer
question_model = T5ForConditionalGeneration.from_pretrained('ramsrigouthamg/t5_squad_v1')
question_tokenizer = T5Tokenizer.from_pretrained('ramsrigouthamg/t5_squad_v1')

def get_question(sentence,answer,mdl,tknizer):
  text = "context: {} answer: {}".format(sentence,answer)
  print (text)
  max_len = 256
  encoding = tknizer.encode_plus(text,max_length=max_len, pad_to_max_length=False,truncation=True, return_tensors="pt")

  input_ids, attention_mask = encoding["input_ids"], encoding["attention_mask"]

  outs = mdl.generate(input_ids=input_ids,
                                  attention_mask=attention_mask,
                                  early_stopping=True,
                                  num_beams=5,
                                  num_return_sequences=1,
                                  no_repeat_ngram_size=2,
                                  max_length=300)


  dec = [tknizer.decode(ids,skip_special_tokens=True) for ids in outs]


  Question = dec[0].replace("question:","")
  Question= Question.strip()
  return Question


context = "Elon Musk said that Tesla will not accept payments in Bitcoin because of environmental concerns."
answer = "Elon Musk"

ques = get_question(context,answer,question_model,question_tokenizer)
print ("question: ",ques)

import gradio as gr

title = "Question Generator Three"
description = "Paste or write a text. Provide a short answer or noun keywords. Submit and the machine will attempt to generate a coherent question"
context = gr.inputs.Textbox(lines=5, placeholder="Enter paragraph/context here...")
answer = gr.inputs.Textbox(lines=3, placeholder="Enter answer/keyword here...")
question = gr.outputs.Textbox( type="auto", label="Question")
examples = [
                                ["""Fears of a new Covid-19 cluster linked to a hotpot restaurant have surfaced amid Hong Kong’s Omicron-fuelled fifth wave, while infections tied to an investment bank continued to expand, triggering the evacuation of residents in a building after vertical transmission of the virus was detected.
On Wednesday, hundreds thronged Covid-19 testing stations in Tuen Mun, with some residents complaining of long waiting times and chaotic arrangements. Authorities have deemed the district a high-risk area because of a higher number of infections.
Health officials said sewage testing would be conducted in Tuen Mun to monitor the spread of the coronavirus, but a string of preliminary-positive cases detected across the city suggested a wider, more worrying situation.
""", "a higher number of infections"]
                               
]

def generate_question(context,answer):
  return get_question(context,answer,question_model,question_tokenizer)

iface = gr.Interface(
  fn=generate_question, 
  inputs=[context,answer], 
  outputs=question, title=title, description=description, examples=examples)
iface.launch(debug=False)