mirror of
https://github.com/BoardWare-Genius/jarvis-models.git
synced 2025-12-13 16:53:24 +00:00
30 lines
1.2 KiB
Python
30 lines
1.2 KiB
Python
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
|