Cristian's picture
code updated
4496ce9
raw
history blame
1.77 kB
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()