# 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()