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)