text-classification / model /sentiment.py
manue
Refactor sentiment model initialization to use AutoModel and AutoTokenizer
e1379c5
raw
history blame
512 Bytes
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
from typing import Any
class Sentiment:
def __init__(self, line: str) -> (list | list[Any] | Any | None):
model = AutoModelForSequenceClassification.from_pretrained("./feel-it-italian-sentiment")
tokenizer = AutoTokenizer.from_pretrained("./feel-it-italian-sentiment")
self.pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
self.result = self.pipe(line)