|
|
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM |
|
from sentence_transformers import SentenceTransformer |
|
|
|
|
|
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)}") |
|
|
|
|
|
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)}") |
|
|
|
|
|
class MetricsModels: |
|
def __init__(self): |
|
self.sentence_bert_model = 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 |
|
|
|
|
|
classifier_model = ClassifierModel() |
|
paraphraser_model = ParaphraserModel() |
|
metrics_models = MetricsModels() |