server.py

import logging
import grpc
from concurrent import futures
import protos.resume_pb2
import protos.resume_pb2_grpc
from huggingface_hub import login
import os
from predictor import Predictor
from transformers import pipeline
from datetime import date

HF_TOKEN = os.environ["HF_Token"]
PORT = os.environ.get("PORT", "50051")
login(HF_TOKEN)

class Resume(protos.resume_pb2_grpc.ResumeServicer):

 def __init__(self):
 self.done = False
 self.logger = logging.getLogger(__name__) 
 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)
 },
 paths=[
 "data/W020230619818476939351.xls",
 "data/W020230619818476975218.xls"
 ],
 today=date(2023,4,1)
 )
 self.done = True

 def Health(self, request, context):
 self.logger.info("Health check")
 if request.ping=="PING":
 if self.done:
 return protos.resume_pb2.PongResponse(done="OK")
 else:
 return protos.resume_pb2.PongResponse(done="Pending")
 return protos.resume_pb2.PongResponse(done="PING request is not valid")

 def GetInfo(self, request, context):
 entities = self.predictor(request.text)
 logging.info(entities)
 res = protos.resume_pb2.ResumeEntitiesResponse()
 for name in entities['name']:
 res.names.append(protos.resume_pb2.Entity(
 entity=name['entity'],
 start=name['start'],
 end=name['end'],
 text=name.get('text',None),
 origin=name["origin"]
 ))
 for age in entities['age']:
 res.ages.append(protos.resume_pb2.Entity(
 entity=age['entity'],
 start=age['start'],
 end=age['end'],
 text=age.get('text',None),
 origin=age["origin"]
 ))
 for gender in entities['gender']:
 res.genders.append(protos.resume_pb2.Entity(
 entity=gender['entity'],
 start=gender['start'],
 end=gender['end'],
 text=gender.get('text',None),
 origin=gender["origin"]
 ))
 for email in entities['email']:
 res.emails.append(protos.resume_pb2.Entity(
 entity=email['entity'],
 start=email['start'],
 end=email['end'],
 text=email.get('text',None),
 origin=email["origin"]
 ))
 for phone in entities['phone']:
 res.phones.append(protos.resume_pb2.Entity(
 entity=phone['entity'],
 start=phone['start'],
 end=phone['end'],
 text=phone.get('text',None),
 origin=phone["origin"]
 ))
 for edu in entities['edus']:
 res.edus.append(protos.resume_pb2.Entity(
 entity=edu['entity'],
 start=edu['start'],
 end=edu['end'],
 text=edu.get('text',None),
 origin=edu["origin"]
 ))
 for school in entities['schools']:
 res.schools.append(protos.resume_pb2.Entity(
 entity=school['entity'],
 start=school['start'],
 end=school['end'],
 text=school.get('text',None),
 origin=school["origin"],
 level=school.get('level',None)
 ))
 for company,start,end in entities['jobs']:
 jobEntity = protos.resume_pb2.JobEntity(
 start=protos.resume_pb2.Entity(
 entity=start['entity'],
 start=start['start'],
 end=start['end'],
 text=start.get('text',None),
 origin=start["origin"]
 ),
 end=protos.resume_pb2.Entity(
 entity=end['entity'],
 start=end['start'],
 end=end['end'],
 text=end.get('text',None),
 origin=end["origin"]
 ),
 company=protos.resume_pb2.Entity(
 entity=company['entity'],
 start=company['start'],
 end=company['end'],
 text=company.get('text',None),
 origin=company["origin"]
 )
 )
 res.jobs.append(jobEntity)
 for title in entities['titles']:
 res.titles.append(protos.resume_pb2.Entity(
 entity=title['entity'],
 start=title['start'],
 end=title['end'],
 text=title.get('text',None),
 origin=title["origin"]
 )) 
 res.work_years = entities['work_time']
 return res
 

def serve(port = "50051"):
 logger = logging.getLogger(__name__)
 server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
 protos.resume_pb2_grpc.add_ResumeServicer_to_server(Resume(), server)
 server.add_insecure_port('[::]:' + port)
 logger.info("Starting server on port %s", port)
 server.start()
 logger.info("Running..")
 server.wait_for_termination()

if __name__ == '__main__':
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 serve(PORT)