from transformers import AutoTokenizer, pipeline
from optimum.intel import OVModelForSequenceClassification 
import util.dict_utils as du

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


    def analyze(self, text: str) -> str:
      if self.pipeline is None:
         self.create_pipeline()
      response = self.pipeline(text)[0]

      return {
        "positive": du.get_item_by_key(response, "label", "Positiv")["score"],
        "neutral": du.get_item_by_key(response, "label", "Neutral")["score"],
        "negative": du.get_item_by_key(response, "label", "Negativ")["score"],
      }

    
    def create_pipeline(self):
      self.pipeline = pipeline('sentiment-analysis', model=self.model, tokenizer=self.tokenizer, top_k=None)
