import os
import sys
sys.path.insert(0, os.path.abspath('./'))

import nltk
nltk.download('punkt')

import numpy as np
from tqdm.auto import tqdm
from transformers import T5Tokenizer
from transformers import AutoModelForSeq2SeqLM
from summarization_dataset import *

import gradio as gr


def predict(data):
    data=[data]
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    """Load Tokenizer"""  
    tokenizer = T5Tokenizer.from_pretrained("t5-small", model_max_length=512)
    
    """Tokenized Inputs"""
    tokenized_inputs = tokenizer(
            data,
            add_special_tokens=True,
            padding="max_length",
            return_token_type_ids=True,
            truncation=True)
    
    ids = torch.tensor(tokenized_inputs['input_ids']).to(device)
    mask = torch.tensor(tokenized_inputs['attention_mask']).to(device)
    
    """Load Model"""
    model_path = "./"


    model = AutoModelForSeq2SeqLM.from_pretrained(model_path + "/summarization_final",device_map=device)
    model.to(device)
    
    
    """Make Prediction"""
    model.eval()

    with torch.no_grad():
        generate_token = model.generate(ids, attention_mask=mask, max_new_tokens=50)
    
    decoded_generate = tokenizer.batch_decode(generate_token.cpu(), skip_special_tokens=True)
    generated = ["\n".join(nltk.sent_tokenize(gen)) for gen in decoded_generate]
    
    return generated[0]


title="Financial News Summarization"
description=""
article="Domain adaptation and fine-tunning dataset from Zhou, Z., Ma, L., & Liu, H. (2021)."
example_list=["EWC provides guests with a first class, professional waxing experience by the most highly trained estheticians in the industry, within the privacy of clean, individual waxing suites. They're so certain everyone will love the EWC experience, European Wax Center offers a free complimentary wax to each new guest. EWC continues to revolutionize the waxing category with their innovative, signatureComfort Wax. This proprietary blend is formulated with the highest quality ingredients to leave skin feeling smooth and make waxing a more pleasant, virtually painless experience.To help enhance and extend waxing services after leaving the center, EWC offers a full collection of proprietary products in the skincare, body, and brow categories."]


# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
                    inputs="text", # what are the inputs?
                    outputs="text", # our fn has two outputs, therefore we have two outputs
                    examples=example_list, 
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch(debug=False, share=False)