Spaces:
Sleeping
Sleeping
Update model_loader.py
Browse files- model_loader.py +1 -14
model_loader.py
CHANGED
@@ -1,7 +1,6 @@
|
|
1 |
# model_loader.py
|
2 |
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM
|
3 |
from sentence_transformers import SentenceTransformer
|
4 |
-
from transformers import pipeline
|
5 |
|
6 |
# Classifier Model (XLM-RoBERTa for toxicity classification)
|
7 |
class ClassifierModel:
|
@@ -39,28 +38,16 @@ class ParaphraserModel:
|
|
39 |
except Exception as e:
|
40 |
raise Exception(f"Error loading paraphrase model or tokenizer: {str(e)}")
|
41 |
|
42 |
-
# Metrics Models (Sentence-BERT
|
43 |
class MetricsModels:
|
44 |
def __init__(self):
|
45 |
self.sentence_bert_model = None
|
46 |
-
self.emotion_classifier = None
|
47 |
-
self.nli_classifier = None
|
48 |
|
49 |
def load_sentence_bert(self):
|
50 |
if self.sentence_bert_model is None:
|
51 |
self.sentence_bert_model = SentenceTransformer('all-MiniLM-L6-v2')
|
52 |
return self.sentence_bert_model
|
53 |
|
54 |
-
def load_emotion_classifier(self):
|
55 |
-
if self.emotion_classifier is None:
|
56 |
-
self.emotion_classifier = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion", top_k=None)
|
57 |
-
return self.emotion_classifier
|
58 |
-
|
59 |
-
def load_nli_classifier(self):
|
60 |
-
if self.nli_classifier is None:
|
61 |
-
self.nli_classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
|
62 |
-
return self.nli_classifier
|
63 |
-
|
64 |
# Singleton instances
|
65 |
classifier_model = ClassifierModel()
|
66 |
paraphraser_model = ParaphraserModel()
|
|
|
1 |
# model_loader.py
|
2 |
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM
|
3 |
from sentence_transformers import SentenceTransformer
|
|
|
4 |
|
5 |
# Classifier Model (XLM-RoBERTa for toxicity classification)
|
6 |
class ClassifierModel:
|
|
|
38 |
except Exception as e:
|
39 |
raise Exception(f"Error loading paraphrase model or tokenizer: {str(e)}")
|
40 |
|
41 |
+
# Metrics Models (Sentence-BERT only)
|
42 |
class MetricsModels:
|
43 |
def __init__(self):
|
44 |
self.sentence_bert_model = None
|
|
|
|
|
45 |
|
46 |
def load_sentence_bert(self):
|
47 |
if self.sentence_bert_model is None:
|
48 |
self.sentence_bert_model = SentenceTransformer('all-MiniLM-L6-v2')
|
49 |
return self.sentence_bert_model
|
50 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
51 |
# Singleton instances
|
52 |
classifier_model = ClassifierModel()
|
53 |
paraphraser_model = ParaphraserModel()
|