Update model_loader.py

model_loader.py

@@ -1,33 +1,67 @@
 # model_loader.py
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-from transformers import AutoModelForCausalLM
-
-def load_classifier_model_and_tokenizer():
-    """
-    Load the fine-tuned XLM-RoBERTa model and tokenizer for toxic comment classification.
-    Returns the model and tokenizer.
-    """
-    try:
-        model_name = "JanviMl/xlm-roberta-toxic-classifier-capstone"
-        model = AutoModelForSequenceClassification.from_pretrained(model_name)
-        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-        return model, tokenizer
-    except Exception as e:
-        raise Exception(f"Error loading classifier model or tokenizer: {str(e)}")
-
-def load_paraphrase_model_and_tokenizer():
-    """
-    Load the Granite 3.2-2B-Instruct model and tokenizer for paraphrasing.
-    Returns the model and tokenizer.
-    """
-    try:
-        model_name = "ibm-granite/granite-3.2-2b-instruct"
-        model = AutoModelForCausalLM.from_pretrained(model_name)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        return model, tokenizer
-    except Exception as e:
-        raise Exception(f"Error loading paraphrase model or tokenizer: {str(e)}")
-
-# Load both models and tokenizers at startup
-classifier_model, classifier_tokenizer = load_classifier_model_and_tokenizer()
-paraphrase_model, paraphrase_tokenizer = load_paraphrase_model_and_tokenizer()
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
+
+# Classifier Model (XLM-RoBERTa for toxicity classification)
+class ClassifierModel:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.load_model()
+
+    def load_model(self):
+        """
+        Load the fine-tuned XLM-RoBERTa model and tokenizer for toxic comment classification.
+        """
+        try:
+            model_name = "JanviMl/xlm-roberta-toxic-classifier-capstone"
+            self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+        except Exception as e:
+            raise Exception(f"Error loading classifier model or tokenizer: {str(e)}")
+
+# Paraphraser Model (Granite 3.2-2B-Instruct for paraphrasing)
+class ParaphraserModel:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.load_model()
+
+    def load_model(self):
+        """
+        Load the Granite 3.2-2B-Instruct model and tokenizer for paraphrasing.
+        """
+        try:
+            model_name = "ibm-granite/granite-3.2-2b-instruct"
+            self.model = AutoModelForCausalLM.from_pretrained(model_name)
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        except Exception as e:
+            raise Exception(f"Error loading paraphrase model or tokenizer: {str(e)}")
+
+# Metrics Models (Sentence-BERT, Emotion Classifier, NLI)
+class MetricsModels:
+    def __init__(self):
+        self.sentence_bert_model = None
+        self.emotion_classifier = None
+        self.nli_classifier = None
+
+    def load_sentence_bert(self):
+        if self.sentence_bert_model is None:
+            self.sentence_bert_model = SentenceTransformer('all-MiniLM-L6-v2')
+        return self.sentence_bert_model
+
+    def load_emotion_classifier(self):
+        if self.emotion_classifier is None:
+            self.emotion_classifier = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion", top_k=None)
+        return self.emotion_classifier
+
+    def load_nli_classifier(self):
+        if self.nli_classifier is None:
+            self.nli_classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+        return self.nli_classifier
+
+# Singleton instances
+classifier_model = ClassifierModel()
+paraphraser_model = ParaphraserModel()
+metrics_models = MetricsModels()