AIIA / ml-service /model.py
RosiYo's picture
feat[bert]: Requirements
3dd6930
raw
history blame
514 Bytes
# model.py
from transformers import AutoModelForSequenceClassification, TFAutoModelForSequenceClassification, AutoTokenize
class Model:
"""A model class to lead the model and tokenizer"""
def __init__(self) -> None:
pass
def load_model():
model = AutoModelForSequenceClassification.from_pretrained('cardiffnlp/twitter-roberta-base-sentiment')
return model
def load_tokenizer():
tokenizer = AutoTokenize.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
return tokenizer