from transformers import AutoTokenizer, pipeline
from optimum.intel import OVModelForTokenClassification 

class NamedEntityRecognizer:
  def __init__(self, model_name):
    self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
    self.model = OVModelForTokenClassification.from_pretrained(model_name, from_pt=True, export=True, load_in_half_precision=True)
    self.create_pipeline()

  
  def recognize(self, text: str):
    if self.pipeline is None:
      self.create_pipeline()

    return self.pipeline(text)
  

  def create_pipeline(self):
    self.pipeline = pipeline('ner', model=self.model, tokenizer=self.tokenizer)