import logging
import os
import grpc
import core.ner as ner
import protobuf.named_entity_recognition_pb2 as ner_pb2

class NamedEntityRecognizerServer():
  def __init__(self) -> None:
    model_name = os.getenv("NER_MODEL", "Davlan/distilbert-base-multilingual-cased-ner-hrl")
    self.recognizer = ner.NamedEntityRecognizer(model_name=model_name)


  async def RecognizeEntities(
    self,
    request: ner_pb2.RecognizeEntitiesRequest,
    context: grpc.aio.ServicerContext,
  ) -> ner_pb2.RecognizeEntitiesResponse:
    response = self.recognizer.recognize(request.text)

    logging.info(f"Recognized entities: {response}")

    entities = []
    for entity in response:
        entities.append(ner_pb2.Entity(
            name=entity["word"],
            type=entity["entity"],
            salience=entity["score"],
        ))

    return ner_pb2.RecognizeEntitiesResponse(entities=entities)