fix: 移除平均池化

app.py

@@ -18,14 +18,6 @@ config = AutoConfig.from_pretrained("Salesforce/codet5p-110m-embedding", trust_r
 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()
-    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
-    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-    return sum_embeddings / sum_mask
-
 @spaces.GPU
 def generate(query1, query2, source_code):
     if len(query1) < 1:
@@ -47,10 +39,10 @@ def generate(query1, query2, source_code):
     # 特殊处理 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].unsqueeze(0), embeddings[2].unsqueeze(0))
-    score2 = cos_sim(embeddings[1].unsqueeze(0), embeddings[2].unsqueeze(0))
+        embeddings = model6(**inputs)[0]  # 直接使用模型输出的嵌入
+
+    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