fix: codet5p 池化张量长度不一致

app.py

@@ -20,9 +20,11 @@ model6 = AutoModel.from_pretrained("Salesforce/codet5p-110m-embedding", config=c
 
 # 创建一个简单的平均池化函数来获取嵌入
 def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
+    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)
+    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):
@@ -47,8 +49,8 @@ def generate(query1, query2, source_code):
     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])
+    score1 = cos_sim(embeddings[0].unsqueeze(0), embeddings[2].unsqueeze(0))
+    score2 = cos_sim(embeddings[1].unsqueeze(0), embeddings[2].unsqueeze(0))
     results.append([model_names[-1], float(score1), float(score2)])
 
     return results