BLAGO_AI / model.py
RomanMaloev's picture
Upload 7 files
91860b8 verified
Raw
History Blame Contribute Delete
459 Bytes
from transformers import AutoModelForSequenceClassification, AutoTokenizer
class BlagoModel:
def __init__(self, model_name="bert-base-uncased"):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
def process(self, text):
inputs = self.tokenizer(text, return_tensors="pt")
outputs = self.model(**inputs)
return outputs.logits