import gradio as gr
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationChain  
from transformers import pipeline

model_name="nateraw/bert-base-uncased-emotion"
model = pipeline('text-classification', model_name, truncation=True)

from transformers import AutoTokenizer, AutoModelWithLMHead
model_name = "mrm8488/t5-base-finetuned-emotion"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model_t5 = AutoModelWithLMHead.from_pretrained(model_name)

def get_emotion(text):
  input_ids = tokenizer.encode(text + '</s>', return_tensors='pt')
  output = model_t5.generate(input_ids=input_ids, return_dict_in_generate=True, output_scores=True)
  transition_scores = model_t5.compute_transition_scores(output.sequences, [s.softmax(dim=1) for s in output.scores], normalize_logits=False)
  dec = [tokenizer.decode(ids) for ids in output.sequences]
  score = transition_scores.min().item()
  return f"{dec[0].replace('<pad>','').replace('</s>','').strip()} [{score}]"

chat = ChatOpenAI()
conversation = ConversationChain(llm=chat)  
#Write a text example of someone angry
with gr.Blocks() as demo:
    label_text = gr.Textbox(label="Sentiment Text", text="")
    chatbot = gr.Chatbot(scale=2)
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])
            

    def respond(message, chat_history):
        bot_message = conversation.run(message)
        chat_history.append((message, bot_message))
        
        l = model(bot_message)[0]
        label_value = f"{l['label']} [{l['score']}]"
        label_value_t5 = get_emotion(bot_message)
        
        return "", chat_history, f"Model1: {label_value_t5} - Model2: {label_value}"

    msg.submit(respond, [msg, chatbot], [msg, chatbot, label_text])

demo.launch()