import gradio as gr
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer, AutoConfig
from scipy.special import softmax
from transformers import pipeline
import torch


## Requirements
model_path = f"eyounge/younge-distilbert-sent-analysis-model"
tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
config = AutoConfig.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

# Preprocess text (username and link placeholders)
def preprocess(Input_text):
    new_text = []
    for t in text.split(" "):
        t = '@user' if t.startswith('@') and len(t) > 1 else t
        t = 'http' if t.startswith('http') else t
        new_text.append(t)
    return " ".join(new_text)

def sentiment_analysis(STATEMENT_ON_COVID_VACCINATION):
    Message = preprocess(STATEMENT_ON_COVID_VACCINATION)

    # PyTorch-based models
    encoded_input = tokenizer(Message, return_tensors='pt')
    output = model(**encoded_input)
    scores_ = output[0][0].detach().numpy()
    scores_ = softmax(scores_)
    
    # Format output dict of scores
    labels = ['Negative', 'Neutral', 'Positive']
    scores = {l:float(s) for (l,s) in zip(labels, scores_) }
    
    return scores

demo = gr.Interface(
    fn=sentiment_analysis, 
    inputs=gr.Textbox(placeholder="Write your tweet here..."), 
    outputs="label", 
    interpretation="default",
    title='SENTIMENT ANALYSIS ON COVID VACCINATION',
    description='Get a sentiment on your input message as Negative/Positive/Neutral'
    allow_flagging=False,
    Caution =[["COVID-19 is real!"]])

demo.launch(inline=False)