Update model_loader.py

model_loader.py

@@ -1,19 +1,33 @@
+# model_loader.py
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import AutoModelForCausalLM
 
-def load_model_and_tokenizer():
+def load_classifier_model_and_tokenizer():
     """
-    Load the fine-tuned XLM-RoBERTa model and tokenizer.
-    Returns the model and tokenizer for use in classification.
+    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"  # Replace with your model repo ID
-        # If the model is local: model_name = "./model"
-        
+        model_name = "JanviMl/xlm-roberta-toxic-classifier-capstone"
         model = AutoModelForSequenceClassification.from_pretrained(model_name)
-        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)  # Use slow tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
         return model, tokenizer
     except Exception as e:
-        raise Exception(f"Error loading model or tokenizer: {str(e)}")
+        raise Exception(f"Error loading classifier model or tokenizer: {str(e)}")
 
-# Load the model and tokenizer once at startup
-model, tokenizer = load_model_and_tokenizer()
+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()