fix: code5t 部署

app.py

@@ -1,8 +1,9 @@
 import spaces
 import gradio as gr
 from numpy.linalg import norm
-from transformers import AutoModel
+from transformers import AutoModel, AutoTokenizer, AutoConfig
 from sentence_transformers import SentenceTransformer
+import torch
 
 cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
 
@@ -11,7 +12,17 @@ model2 = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-en", trust_re
 model3 = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-zh", trust_remote_code=True)
 model4 = SentenceTransformer("aspire/acge_text_embedding")
 model5 = SentenceTransformer("intfloat/multilingual-e5-large")
-model6 = SentenceTransformer("Salesforce/codet5p-110m-embedding")
+
+# 对于 Salesforce/codet5p-110m-embedding 模型，我们需要特殊处理
+config = AutoConfig.from_pretrained("Salesforce/codet5p-110m-embedding", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5p-110m-embedding", trust_remote_code=True)
+model6 = AutoModel.from_pretrained("Salesforce/codet5p-110m-embedding", config=config, trust_remote_code=True)
+
+# 创建一个简单的平均池化函数来获取嵌入
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
 @spaces.GPU
 def generate(query1, query2, source_code):
@@ -24,11 +35,21 @@ def generate(query1, query2, source_code):
 
     results = []
     model_names = ["jinaai/jina-embeddings-v2-base-code", "jinaai/jina-embeddings-v2-base-en", "jinaai/jina-embeddings-v2-base-zh", "aspire/acge_text_embedding", "intfloat/multilingual-e5-large", "Salesforce/codet5p-110m-embedding"]
-    for model, name in zip([model1, model2, model3, model4, model5, model6], model_names):
+    
+    for model, name in zip([model1, model2, model3, model4, model5], model_names[:-1]):
         embeddings = model.encode([query1, query2, source_code])
         score1 = cos_sim(embeddings[0], embeddings[2])
         score2 = cos_sim(embeddings[1], embeddings[2])
         results.append([name, float(score1), float(score2)])
+    
+    # 特殊处理 Salesforce/codet5p-110m-embedding 模型
+    inputs = tokenizer([query1, query2, source_code], padding=True, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        model_output = model6(**inputs)
+    embeddings = mean_pooling(model_output, inputs['attention_mask'])
+    score1 = cos_sim(embeddings[0], embeddings[2])
+    score2 = cos_sim(embeddings[1], embeddings[2])
+    results.append([model_names[-1], float(score1), float(score2)])
 
     return results