#importing the spacy and bert model
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import pipeline
import gradio as gr
import re
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('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):
  text = re.sub(r'Photo by.+', '', text)
  text = re.sub(r"\n", " ", text)
  text = re.sub(r"\n\n", " ", text)
  text = re.sub(r"\t", " ", text)
  text = text.strip(" ")
  text = re.sub(
      " +", " ", text
      ).strip()  # get rid of multiple spaces and replace with a single
      
  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):
  text = re.sub(r'Photo by.+', '', text)
  text = re.sub(r"\n", " ", text)
  text = re.sub(r"\n\n", " ", text)
  text = re.sub(r"\t", " ", text)
  text = text.strip(" ")
  text = re.sub(
          " +", " ", text
      ).strip()  # get rid of multiple spaces and replace with a single
      
  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 = 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)