import gradio
from transformers import pipeline
import nltk

nltk.download('punkt')

summarizer = pipeline('summarization', model='t5-base')

classifier_model_name = 'manvik28/FinBERT_Tuned'
classifier_emotions = ['positive', 'neutral', 'negative']

classifier = pipeline('text-classification', model=classifier_model_name)  

def find_emotional_sentences(text, emotions, threshold):
    sentences_by_emotion = {}
    for e in emotions:
        sentences_by_emotion[e]=[]
    sentences = nltk.sent_tokenize(text)
    print(f'Document has {len(text)} characters and {len(sentences)} sentences.')
    for s in sentences:
        prediction = classifier(s)
        if (prediction[0]['label']!='neutral' and prediction[0]['score']>threshold):
            #print (f'Sentence #{sentences.index(s)}: {prediction} {s}')
            sentences_by_emotion[prediction[0]['label']].append(s)
    for e in emotions:
        print(f'{e}: {len(sentences_by_emotion[e])} sentences')
    return sentences_by_emotion

def summarize_sentences(sentences_by_emotion, min_length, max_length):
    summary = {}
    for k in sentences_by_emotion.keys():
        if (len(sentences_by_emotion[k])!=0):
            text = ' '.join(sentences_by_emotion[k])
            summary = summarizer(text, min_length=min_length, max_length=max_length)
            print(f"{k.upper()}: {summary[0]['summary_text']}\n")

    return  summary
    
def my_inference_function(sec_txt):
    sentences_by_emotion = find_emotional_sentences(sec_txt, classifier_emotions, 0.85)
    summarized_text = summarize_sentences(sentences_by_emotion, 50, 150)
    return summarized_text
    
gradio_interface = gradio.Interface(
  fn = my_inference_function,
  inputs = "text",
  outputs = "text"
)
gradio_interface.launch()