import os
import gradio
from PIL import Image
from timeit import default_timer as timer
from tensorflow import keras
import torch
from transformers import AutoTokenizer, TFAutoModelForSeq2SeqLM, create_optimizer, DataCollatorForSeq2Seq
import numpy as np

loaded_model = TFAutoModelForSeq2SeqLM.from_pretrained("runaksh/financial_summary_T5_base")
loaded_tokenizer = AutoTokenizer.from_pretrained("runaksh/financial_summary_T5_base")

# Function for generating summary
def generate_summary(text,min_length=55,max_length=80):
  text = "summarize: "+text
  input = loaded_tokenizer(text,max_length=512,truncation=True,return_tensors='tf').input_ids
  op=loaded_model.generate(input,min_length=min_length,max_length=max_length)
  decoded_op = loaded_tokenizer.batch_decode(op,skip_special_tokens=True)
  return decoded_op

title = "Financial News Summary"
description = "Enter the news"

# Gradio elements

# Input from user
in_prompt = gradio.components.Textbox(lines=2, label='Enter the News')

# Output response
out_response = gradio.components.Textbox(label='Summary')

# Gradio interface to generate UI link
iface = gradio.Interface(fn=generate_summary,
                         inputs = in_prompt,
                         outputs = out_response,
                         title=title,
                         description=description
                         )

iface.launch(debug = True)