import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration
from transformers.optimization import  Adafactor 
import time
import warnings
tokenizer = T5Tokenizer.from_pretrained('t5-base')
model = T5ForConditionalGeneration.from_pretrained('pytoch_model.bin', return_dict=True,config='t5-base-config.json')

def generate(text):
  model.eval()
  input_ids = tokenizer.encode("WebNLG:{} </s>".format(text), return_tensors="pt")  # Batch size 1
  # input_ids.to(dev)
  s = time.time()
  outputs = model.generate(input_ids)
  gen_text=tokenizer.decode(outputs[0]).replace('<pad>','').replace('</s>','')
  elapsed = time.time() - s
  print('Generated in {} seconds'.format(str(elapsed)[:4]))

  
  return gen_text


import gradio as gr



# Define the Gradio interface
iface = gr.Interface(
    fn=generate,  # Replace with your actual function
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter text here..."),
    outputs=gr.outputs.Textbox(),
    title="Text Generation App",
    description="Enter some data (Example : Russia | leader | Putin)",
)

# Launch the Gradio interface
iface.launch()