Update app.py

app.py

@@ -3,49 +3,44 @@ 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)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     trust_remote_code=True,
-    device_map="auto",
     torch_dtype=torch.bfloat16
-)
+).to(device)
+
 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"]
+    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=True).to(device)
     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]
+        embeddings = embedding_layer(inputs["input_ids"])
+        sentence_embedding = embeddings.mean(dim=1)
     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="我爱上海")
+        gr.Textbox(label="Sentence 2", placeholder="我爱上海"),
     ],
     outputs=gr.Textbox(label="Similarity"),
-    title=title,
+    title="Calculate two sentences's similarity by ERNIE 4.5-0.3B's embedding layer",
     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()
+    demo.launch(share=True)
+