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)