File size: 490 Bytes
cd60e9d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
from transformers import PretrainedConfig
mp = {0:'sad',1:'joy',2:'love',3:'anger',4:'fear',5:'surprise'}
class SentimentConfig(PretrainedConfig):
model_type = "SententenceTransformerSentimentClassifier"
def __init__(self, embedding_model: str="all-MiniLM-L6-v2", class_map: dict=mp, h1: int=44, h2: int=46, **kwargs):
self.embedding_model = embedding_model
self.class_map = class_map
self.h1 = h1
self.h2 = h2
super().__init__(**kwargs) |