#importing the spacy and bert model
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import pipeline
import gradio as gr
from gradio.mix import Parallel
import spacy
import pandas as pd

#Intializing the spacy model for NER and the finbert model for sentiment analysis
nlp = spacy.load('spacy/en_core_web_sm')
finbert = BertForSequenceClassification.from_pretrained('yiyanghkust/finbert-tone',num_labels=3)
tokenizer = BertTokenizer.from_pretrained('yiyanghkust/finbert-tone')
sentiment = pipeline("sentiment-analysis", model=finbert, tokenizer=tokenizer)


#defining a function to give us the sentiment of the article
def return_sentiment(text):
  results = sentiment(text[:512])
  return (f"{results[0]['label']} ---> {results[0]['score']}")
  
#defining a function to return the names of the organization present in the article
def show_org(text):
  org = []
  doc = nlp(text)
  if doc.ents:
    for ent in doc.ents:
      if ent.label_ == 'ORG':
        org.append(ent.text)
  None
  
  df = pd.DataFrame()
  org = list(set(org))
  df['Organization'] = org
  return df
  
  
sentiment_analysis = gr.Interface(
    fn=return_sentiment,
    inputs = gr.inputs.Textbox(label="Input your news article here", optional=False),
    outputs=gr.outputs.Textbox(label="Sentiment Analysis"),
)

named_organization = gr.Interface(
    fn=show_org,
    inputs = gr.inputs.Textbox(label="Input your news article here", optional=False),
    outputs=gr.outputs.Dataframe(label="Named organizations"),
)

Parallel(
    sentiment_analysis,
    named_organization,
    title="Sentiment Analysis of stock news",
    inputs=gr.inputs.Textbox(
        label="Input your article here",
    ),
    theme="darkhuggingface",
).launch(enable_queue=True, debug=True)