nlp / app.py
hvtham's picture
Update app.py
6cb3a35
pip install transformers
pip install gradio
import transformers
from transformers import AutoTokenizer, AutoModel
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
import torch
import torchvision
import gradio as gr
import io
from PIL import Image
def similarity_score(text1, text2):
# Load PhoBERT tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base")
model = AutoModel.from_pretrained("vinai/phobert-base", output_hidden_states=True)
# Encode the two texts using PhoBERT tokenizer
tokens1 = tokenizer(text1, return_tensors='pt', padding=True)
tokens2 = tokenizer(text2, return_tensors='pt', padding=True)
# Pass the encoded tokens to PhoBERT to get the hidden states
with torch.no_grad():
hidden_states1 = model(tokens1['input_ids'], tokens1['attention_mask'])
hidden_states2 = model(tokens2['input_ids'], tokens2['attention_mask'])
# Calculate cosine similarity between the last hidden states of the two texts
similarity = cosine_similarity(hidden_states1[0][:, -1, :], hidden_states2[0][:, -1, :])[0][0]
#return similarity
pca = PCA(n_components=2)
pca.fit_transform(torch.cat((hidden_states1[0][:, -1, :], hidden_states2[0][:, -1, :])))
# tạo biểu đồ
fig, ax = plt.subplots()
ax.scatter(pca.components_[0][0], pca.components_[1][0], color='red')
ax.text(pca.components_[0][0], pca.components_[1][0], text1, fontsize=12)
ax.scatter(pca.components_[0][1], pca.components_[1][1], color='blue')
ax.text(pca.components_[0][1], pca.components_[1][1], text2, fontsize=12)
ax.set_title(f"Khoảng các Consin của 2 câu: {similarity:.4f}")
# chuyển đổi hình ảnh matplotlib sang định dạng bytes và hiển thị trên Gradio
buffer = io.BytesIO()
fig.canvas.print_png(buffer)
img_bytes = buffer.getvalue()
return Image.open(io.BytesIO(img_bytes))
input_text1 = gr.inputs.Textbox(label="Câu 1")
input_text2 = gr.inputs.Textbox(label="Câu 2")
output_score = gr.outputs.Textbox(label="Điểm tương đồng")
app = gr.Interface(
fn=similarity_score,
inputs=[input_text1, input_text2],
outputs=gr.outputs.Image(type='numpy'),
title="So sánh độ giống nhau giữa 2 câu",
description="Nhập vào 2 câu, kết quả trả ra độ tương đồng của 2 câu đó",
allow_flagging=False,
layout="vertical",
theme="default",)
app.launch(share=True)