from transformers import XLNetForSequenceClassification, XLNetTokenizer,BertForSequenceClassification,BertTokenizer, RobertaForSequenceClassification,RobertaTokenizer
import torch
from typing import Dict
import gradio as gr


model = BertForSequenceClassification.from_pretrained("./Personality_detection_Classification_Save/", num_labels=5)#=num_labels)
tokenizer = BertTokenizer.from_pretrained('./Personality_detection_Classification_Save/', do_lower_case=True) 
model.config.label2id= {
"Extroversion": 0,
"Neuroticism": 1,
"Agreeableness": 2,
"Conscientiousness": 3,
"Openness": 4,
}

model.config.id2label={
    "0": "Extroversion",
    "1": "Neuroticism",
    "2": "Agreeableness",
    "3": "Conscientiousness",
    "4": "Openness",}

def Personality_Detection_from_reviews_submitted (model_input: str) -> Dict[str, float]:
    # Encoding input data
    dict_custom={}
    Preprocess_part1=model_input[:len(model_input)]
    Preprocess_part2=model_input[len(model_input):]
    dict1=tokenizer.encode_plus(Preprocess_part1,max_length=1024,padding=True,truncation=True)
    dict2=tokenizer.encode_plus(Preprocess_part2,max_length=1024,padding=True,truncation=True)
    dict_custom['input_ids']=[dict1['input_ids'],dict1['input_ids']]
    dict_custom['token_type_ids']=[dict1['token_type_ids'],dict1['token_type_ids']]
    dict_custom['attention_mask']=[dict1['attention_mask'],dict1['attention_mask']]
    outs = model(torch.tensor(dict_custom['input_ids']), token_type_ids=None, attention_mask=torch.tensor(dict_custom['attention_mask']))
    b_logit_pred = outs[0]
    pred_label = torch.sigmoid(b_logit_pred)
    ret ={
        "Extroversion": float(pred_label[0][0]),
        "Neuroticism": float(pred_label[0][1]),
        "Agreeableness": float(pred_label[0][2]),
        "Conscientiousness": float(pred_label[0][3]),
        "Openness": float(pred_label[0][4]),}
    return ret
model_input = gr.Textbox("Input text here (Note: This model is trained to classify Essays(Still in Progress phase))", show_label=False)
model_output = gr.Label(" Big-Five personality traits Result", num_top_classes=6, show_label=True, label="Big-Five personality traits Labels assigned to this text")
examples = [
    (  "Well, here we go with the stream of consciousness essay. I used to do things like this in high school sometimes.",
       "They were pretty interesting, but I often find myself with a lack of things to say. ",
       "I normally consider myself someone who gets straight to the point. I wonder if I should hit enter any time to send this back to the front",
       "Maybe I'll fix it later. My friend is playing guitar in my room now. Sort of playing anyway.",
       "More like messing with it. He's still learning. There's a drawing on the wall next to me. "
    ),
    (   "An open keyboard and buttons to push. The thing finally worked and I need not use periods, commas and all those thinks.",
        "Double space after a period. We can't help it. I put spaces between my words and I do my happy little assignment of jibber-jabber.",
        "Babble babble babble for 20 relaxing minutes and I feel silly and grammatically incorrect. I am linked to an unknown reader.",
        "A graduate student with an absurd job. I type. I jabber and I think about dinoflagellates. About sunflower crosses and about ",
        "the fiberglass that has be added to my lips via clove cigarettes and I think about things that I shouldn't be thinking.",
        "I know I shouldn't be thinking. or writing let's say/  So I don't. Thoughts don't solidify. They lodge in the back. behind my tongue maybe.",
    )
]

title = "Big Five Personality Traits Detection From Reviews Submitted"
description = "The traditional machine learning models give a lot of pain when we do not have sufficient labeled data for the specific task or domain we care about to train a reliable model. Transfer learning allows us to deal with these scenarios by leveraging the already existing labeled data of some related task or domain. We try to store this knowledge gained in solving the source task in the source domain and apply it to our problem of interest. In this work, I have utilized Transfer Learning utilizing BERT BASE UNCASED model to fine tune on Big-Five Personality traits Dataset."
text1 = (
    #"<center> Author: Owais Ahmad Data Scientist at <b> Thoucentric </b> <a href=\"https://www.linkedin.com/in/owaiskhan9654/\">Visit Profile</a> <br></center>"
    
    # "<center> Model Trained GOOGLE COLAB Kernel <a href=\"https://colab.research.google.com/drive/1dbr9t9vxF8M80FAH9WoBFEdT0iGaXr1r?usp=sharing\">Link</a> <br></center>"
        
    # "<center> Kaggle Profile <a href=\"https://www.kaggle.com/owaiskhan9654\">Link</a> <br> </center>"
    
    "<center> HuggingFace Model Deployed Repository <a href=\"\"> Yet to be deployed</a> <br></center>"
)

app = gr.Interface(
    Personality_Detection_from_reviews_submitted, 
	inputs=model_input,
    outputs=model_output, 
    examples=examples,
    title=title,
	description=description,
    article=text1,
    allow_flagging='never',
    analytics_enabled=False,
)


app.launch(inline=True,share=True, show_error=True)