fix(predictor): fix some error

docker-compose.yml

@@ -8,4 +8,5 @@ services:
       - "50050:50051"
     environment:
       - HF_Token=${HF_Token}
+      - DEVICE=cpu
 

predictor/__init__.py

@@ -46,7 +46,7 @@ class Predictor():
                     + r"空,曾,毋,沙,乜,养,鞠,须,丰,巢,关,蒯,相,查,后,荆,红,游,竺,权,逑,盖,益,桓,公,万俟,司马,上官,欧阳,夏侯,诸葛,闻人,东方,赫连,皇甫,尉迟," \
                     + r"公羊,澹台,公冶,宗政,濮阳,淳于,单于,太叔,申屠,公孙,仲孙,轩辕,令狐,锺离,宇文,长孙,慕容,鲜于,闾丘,司徒,司空,丌官,司寇,仉,督,子车," \
                     + r"颛孙,端木,巫马,公西,漆雕,乐正,壤驷,公良,拓拔,夹谷,宰父,谷梁,晋,楚,阎,法,汝,鄢,涂,钦,段干,百里,东郭,南门,呼延,归,海,羊舌,微生,岳," \
-                    + r"帅,缑,亢,况,后,有,琴,梁丘,左丘,东门,西门,商,牟,佘,佴,伯,赏,南宫,墨,哈,谯,笪,年,爱,阳,佟,第五,言,福,邱]"
+                    + r"帅,缑,亢,况,后,有,琴,梁丘,左丘,东门,西门,商,牟,佘,佴,伯,赏,南宫,墨,哈,谯,笪,年,爱,阳,佟,第五,言,福,邱,钟]"
         first_name = r' {0,3}[\u4e00-\u9fa5]( {0,3}[\u4e00-\u9fa5]){0,3}'
         self.name_pattern = re.compile(last_name + first_name)
         self.phone_pattern = re.compile(r'1 {0,4}(3 {0,4}\d|4 {0,4}[5-9]|5 {0,4}[0-35-9]|6 {0,4}[2567]|7 {0,4}[0-8]|8 {0,4}\d|9 {0,4}[0-35-9]) {0,4}(\d {0,4}){8}')
@@ -106,14 +106,14 @@ class Predictor():
     def to_date(self, datestr:str):
         if re.match("^\d{4}$",datestr):
             return date(int(datestr),1,1)
-        match = re.match("^\d{4}(\D)\d{1,2}",datestr)
+        match = re.match("^(\d{4})\D(\d{1,2})",datestr)
         if match is not None:
             try:
-                m = min(max(int(datestr.split(match.group(1))[1]),1),12)
-                return date(int(datestr.split(match.group(1))[0]),m,1)
+                y = int(match.group(1))
+                m = min(max(int(match.group(2)),1),12)
+                return date(y,m,1)
             except ValueError:
-                print(int(datestr.split(match.group(1))[0]),int(datestr.split(match.group(1))[1]))
-                raise
+                print(datestr)
         if datestr=="至今":
             return self.today
         return None
@@ -206,25 +206,44 @@ class Predictor():
         # 获取名字，先过滤所有空白字符，防止名字中间有空格
         remove_blanks_text, index_mapper = self.remove_blanks(text, re.compile(r' '))
         start_time = time.perf_counter()
+        backup_name = []
         for block in self.split_to_blocks(remove_blanks_text):
             block_text,block_l = block['text'],block['start']
             entities = self.pipelines['name'](block_text)
             for entity in entities:
-                if entity['entity']=='NAME' and self.name_pattern.match(entity['word']) is not None:
-                    obj = {
-                        'start': index_mapper[block_l+entity['start']],
-                        'end': index_mapper[block_l+entity['end']-1]+1,
-                        'entity': 'NAME',
-                        'text': entity['word']
-                    }
-                    repeat = False
-                    for o in return_obj['name']:
-                        if obj['start']==o['start'] and obj['end']==o['end']:
-                            repeat = True
-                            break
-                    if not repeat:
-                        obj['origin'] = text[obj['start']:obj['end']]
-                        return_obj['name'].append(obj)
+                if entity['entity']=='NAME':
+                    if self.name_pattern.match(entity['word']) is not None:
+                        obj = {
+                            'start': index_mapper[block_l+entity['start']],
+                            'end': index_mapper[block_l+entity['end']-1]+1,
+                            'entity': 'NAME',
+                            'text': entity['word']
+                        }
+                        repeat = False
+                        for o in return_obj['name']:
+                            if obj['start']==o['start'] and obj['end']==o['end']:
+                                repeat = True
+                                break
+                        if not repeat:
+                            obj['origin'] = text[obj['start']:obj['end']]
+                            return_obj['name'].append(obj)
+                    else:
+                        obj = {
+                            'start': index_mapper[block_l+entity['start']],
+                            'end': index_mapper[block_l+entity['end']-1]+1,
+                            'entity': 'NAME',
+                            'text': entity['word']
+                        }
+                        repeat = False
+                        for o in return_obj['name']:
+                            if obj['start']==o['start'] and obj['end']==o['end']:
+                                repeat = True
+                                break
+                        if not repeat:
+                            obj['origin'] = text[obj['start']:obj['end']]
+                            backup_name.append(obj)
+        if len(return_obj['name'])==0:
+            return_obj['name'] = backup_name
         end_time = time.perf_counter()
         self.logger.info(f"process name time: {end_time-start_time}")
         # 获取年龄
@@ -301,6 +320,8 @@ class Predictor():
                             break
                     if not repeat:
                         obj['origin'] = text[obj['start']:obj['end']]
+                        if "text" not in obj:
+                            obj['text'] = obj['origin'].split("\n")[-1]
                         return_obj['schools'].append(obj)
         # 正则找学校
         for school_match in re.finditer(r"|".join(self.schools.keys()), remove_blanks_text):
@@ -309,7 +330,7 @@ class Predictor():
                 'start': index_mapper[start],
                 'end': index_mapper[end-1]+1,
                 'entity': 'SCHOOL',
-                'text': school_match.group(),
+                'text': school_match.group().split('\n')[-1],
             }
             repeat = False
             for o in return_obj['schools']:
@@ -320,6 +341,7 @@ class Predictor():
                 obj['origin'] = text[obj['start']:obj['end']]
                 obj['level'] = self.schools[obj['text']]
                 return_obj['schools'].append(obj)
+        return_obj['schools'] = sorted(return_obj['schools'], key=lambda x: x['start'])
         end_time = time.perf_counter()
         self.logger.info(f"process school time: {end_time-start_time}")
         start_time = time.perf_counter()
@@ -410,7 +432,7 @@ class Predictor():
                 diff_m = end.month-start.month
                 work_month += diff_y * 12 + diff_m
                 last_end = end
-            return_obj['work_time'] = math.ceil(work_month/12)
+            return_obj['work_time'] = max(math.ceil(work_month/12),0)
         end_time = time.perf_counter()
         self.logger.info(f"process work time: {end_time-start_time}")
         start_time = time.perf_counter()

server.py

@@ -11,6 +11,7 @@ from datetime import date
 
 HF_TOKEN = os.environ["HF_Token"]
 PORT = os.environ.get("PORT", "50051")
+DEVICE = os.environ.get("DEVICE", "cpu")
 login(HF_TOKEN)
 
 class Resume(protos.resume_pb2_grpc.ResumeServicer):
@@ -22,15 +23,27 @@ class Resume(protos.resume_pb2_grpc.ResumeServicer):
             pipeline=pipeline(
                 "textencode",
                 model="minskiter/cossim-bert-chinese-wwm-ext",
-                device="cpu",
+                device=DEVICE,
                 trust_remote_code=True,
                 use_auth_token=True
             )
         )
         self.predictor = Predictor(
             pipelines={
-                "name": pipeline("nerpipe", model="minskiter/resume-token-classification-name-0708",trust_remote_code=True,use_auth_token=True),
-                "common": pipeline("nerpipe",model="minskiter/resume-token-classification",trust_remote_code=True,use_auth_token=True)
+                "name": pipeline(
+                    "nerpipe", 
+                    device=DEVICE,
+                    model="minskiter/resume-token-classification-name-0708",
+                    trust_remote_code=True,
+                    use_auth_token=True
+                ),
+                "common": pipeline(
+                    "nerpipe",
+                    model="minskiter/resume-token-classification",
+                    device=DEVICE,
+                    trust_remote_code=True,
+                    use_auth_token=True
+                )
             },
             paths=[
                 "data/W020230619818476939351.xls",