fix(server.py): optimize similarity algorithm

predictor/__init__.py

@@ -9,6 +9,8 @@ import queue
 from datetime import date
 import time
 import logging
+import torch
+import torch.nn.functional as F
 
 class Predictor():
 
@@ -443,7 +445,7 @@ class Predictor():
                 'start': start,
                 'end': end,
                 'entity': 'GENDER',
-                'word': text[start:end],
+                'origin': text[start:end],
                 'text': text[start:end]
             })
         end_time = time.perf_counter()
@@ -505,17 +507,22 @@ class PositionPredictor():
         ) -> List[Dict[str, Union[str, float]]]:
             ans = []
             resume_blocks = self.__split_blocks(resume)
-            print(resume_blocks,positions)
+            resume_encoding = []
+            for block_resume in resume_blocks:
+                resume_encoding.append(torch.tensor(self.pipeline(block_resume)[0]))
+            resume_encoding = torch.stack(resume_encoding,dim=0)
             for position in positions:
                 requireds = position['required']
                 score = 0.0
+                block_encodings = []
                 for required in requireds:
                     blocks = self.__split_blocks(required)
-                    print(blocks)
                     for block in blocks:
-                        for block_resume in resume_blocks:
-                            score = max(score, self.pipeline((block_resume, block))[0])
-                            self.logger.info(f"position: {position['name']}, required: {block}, resume: {block_resume}, score: {score}")
+                        block_encodings.append(torch.tensor(self.pipeline(block)[0]))
+                block_encodings = torch.stack(block_encodings,dim=0)
+                cos_sims = F.cosine_similarity(resume_encoding.unsqueeze(1), block_encodings.unsqueeze(0),dim=-1)
+                score = cos_sims.max().item()
+                self.logger.info(f"position: {position['name']}, score: {score}")
                 ans.append({
                     'position': position['name'],
                     'score': score

server.py

@@ -20,8 +20,8 @@ class Resume(protos.resume_pb2_grpc.ResumeServicer):
         self.logger = logging.getLogger(__name__)  
         self.position_predictor =  PositionPredictor(
             pipeline=pipeline(
-                "sentences_sim",
-                model="minskiter/simbert-chinese-bert-wwm-ext",
+                "textencode",
+                model="minskiter/cossim-bert-chinese-wwm-ext",
                 device="cpu",
                 trust_remote_code=True,
                 use_auth_token=True