JanviMl commited on
Commit
0d6bc62
·
verified ·
1 Parent(s): 5a980db

Update model_loader.py

Browse files
Files changed (1) hide show
  1. 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, Emotion Classifier, NLI)
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()