feat: sentiment engine

sentiment_engine/sentiment_engine.py

@@ -0,0 +1,29 @@
+import logging
+
+import onnxruntime
+from transformers import BertTokenizer
+import numpy as np
+
+
+class SentimentEngine():
+
+    def __init__(self, model_path="resources/sentiment_engine/models/paimon_sentiment.onnx"):
+        logging.info('Initializing Sentiment Engine...')
+        onnx_model_path = model_path
+        self.ort_session = onnxruntime.InferenceSession(onnx_model_path, providers=['CPUExecutionProvider'])
+        self.tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
+
+    def infer(self, text):
+        tokens = self.tokenizer(text, return_tensors="np")
+        input_dict = {
+            "input_ids": tokens["input_ids"],
+            "attention_mask": tokens["attention_mask"],
+        }
+        # Convert input_ids and attention_mask to int64
+        input_dict["input_ids"] = input_dict["input_ids"].astype(np.int64)
+        input_dict["attention_mask"] = input_dict["attention_mask"].astype(np.int64)
+        logits = self.ort_session.run(["logits"], input_dict)[0]
+        probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)
+        predicted = np.argmax(probabilities, axis=1)[0]
+        logging.info(f'Sentiment Engine Infer: {predicted}')
+        return predicted