Update model_loader.py

model_loader.py

@@ -1,54 +1,80 @@
 # model_loader.py
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
+import torch
+import os
 
-# Classifier Model (XLM-RoBERTa for toxicity classification)
 class ClassifierModel:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-        self.load_model()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.load_classifier_model()
 
-    def load_model(self):
+    def load_classifier_model(self):
         """
-        Load the fine-tuned XLM-RoBERTa model and tokenizer for toxic comment classification.
+        Load the fine-tuned XLM-RoBERTa model and tokenizer for toxicity classification.
         """
         try:
             model_name = "JanviMl/xlm-roberta-toxic-classifier-capstone"
+            print(f"Loading classifier model: {model_name}")
             self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model.to(self.device)
+            self.model.eval()
+            print("Classifier model loaded successfully")
         except Exception as e:
-            raise Exception(f"Error loading classifier model or tokenizer: {str(e)}")
+            print(f"Error loading classifier model: {str(e)}")
+            raise
+
+classifier_model = ClassifierModel()
 
-# Paraphraser Model (Granite 3.2-2B-Instruct for paraphrasing)
 class ParaphraserModel:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-        self.load_model()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.load_paraphraser_model()
 
-    def load_model(self):
+    def load_paraphraser_model(self):
         """
-        Load the Granite 3.2-2B-Instruct model and tokenizer for paraphrasing.
+        Load the fine-tuned Granite 3.2-2B-Instruct model and tokenizer for paraphrasing.
         """
         try:
             model_name = "ibm-granite/granite-3.2-2b-instruct"
+            print(f"Loading paraphraser model: {model_name}")
             self.model = AutoModelForCausalLM.from_pretrained(model_name)
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            # Set a distinct pad token to avoid conflict with eos token
+            if self.tokenizer.pad_token is None or self.tokenizer.pad_token == self.tokenizer.eos_token:
+                self.tokenizer.pad_token = "<pad>"
+                self.model.config.pad_token_id = self.tokenizer.convert_tokens_to_ids("<pad>")
+            self.model.to(self.device)
+            self.model.eval()
+            print("Paraphraser model loaded successfully")
         except Exception as e:
-            raise Exception(f"Error loading paraphrase model or tokenizer: {str(e)}")
+            print(f"Error loading paraphraser model: {str(e)}")
+            raise
+
+paraphraser_model = ParaphraserModel()
 
-# Metrics Models (Sentence-BERT only)
 class MetricsModels:
     def __init__(self):
-        self.sentence_bert_model = None
+        self.sentence_bert = None
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.load_sentence_bert()
 
     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
+        """
+        Load the Sentence-BERT model for computing semantic similarity.
+        """
+        try:
+            model_name = "sentence-transformers/all-MiniLM-L6-v2"
+            print(f"Loading Sentence-BERT model: {model_name}")
+            self.sentence_bert = SentenceTransformer(model_name, device=self.device)
+            print("Sentence-BERT model loaded successfully")
+        except Exception as e:
+            print(f"Error loading Sentence-BERT model: {str(e)}")
+            raise
 
-# Singleton instances
-classifier_model = ClassifierModel()
-paraphraser_model = ParaphraserModel()
 metrics_models = MetricsModels()