File size: 1,747 Bytes
4d42cbc 75b96db 4d42cbc 75b96db 4d42cbc 75b96db 4d42cbc 75b96db 4d42cbc 75b96db 4d42cbc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
# 模型名称
model_name = "baidu/ERNIE-4.5-0.3B-PT"
# 加载 tokenizer 和模型（首次运行可能较慢）
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.float32,
device_map="auto"
)
embedding_layer = model.get_input_embeddings()
# 提取句子的平均 embedding
def get_sentence_embedding(text):
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=True)
input_ids = inputs["input_ids"]
with torch.no_grad():
embeddings = embedding_layer(input_ids) # shape: [1, seq_len, hidden_size]
sentence_embedding = embeddings.mean(dim=1) # shape: [1, hidden_size]
return sentence_embedding
# Gradio 回调函数
def calculate_similarity(sentence1, sentence2):
emb1 = get_sentence_embedding(sentence1)
emb2 = get_sentence_embedding(sentence2)
similarity = F.cosine_similarity(emb1, emb2).item()
return f"Similarity: {similarity:.4f}"
# Gradio 界面
title = "Calculate two sentences's similarity by ERNIE 4.5-0.3B's embedding layer"
demo = gr.Interface(
fn=calculate_similarity,
inputs=[
gr.Textbox(label="Sentence 1", placeholder="我爱北京"),
gr.Textbox(label="Sentence 2", placeholder="我爱上海")
],
outputs=gr.Textbox(label="Similarity"),
title=title,
description="This app uses the embedding layer of Baidu ERNIE-4.5-0.3B-PT model to compute the cosine similarity between two sentences.",
)
# 启动 Gradio app
if __name__ == "__main__":
demo.launch()
|