Update dataset_utils.py

dataset_utils.py

@@ -1,12 +1,20 @@
 import pandas as pd
 import torch
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset
 from sklearn.preprocessing import LabelEncoder
 from transformers import BertTokenizer, RobertaTokenizer, DebertaTokenizer
 import pickle
 import os
 
-from config import TEXT_COLUMN, LABEL_COLUMNS, MAX_LEN, TOKENIZER_PATH, LABEL_ENCODERS_PATH, METADATA_COLUMNS
+from config import (
+    TEXT_COLUMN,
+    LABEL_COLUMNS,
+    MAX_LEN,
+    TOKENIZER_PATH,
+    LABEL_ENCODERS_PATH,
+    METADATA_COLUMNS,
+    MODEL_NAME  # ✅ Add this in your config.py: MODEL_NAME = "roberta-base"
+)
 
 class ComplianceDataset(Dataset):
     def __init__(self, texts, labels, tokenizer, max_len):
@@ -69,13 +77,14 @@ def load_and_preprocess_data(data_path):
         data[col] = label_encoders[col].fit_transform(data[col])
     return data, label_encoders
 
-def get_tokenizer(model_name):
-    if "bert" in model_name.lower():
-        return BertTokenizer.from_pretrained(model_name)
+def get_tokenizer(model_name=MODEL_NAME):
+    model_name = model_name or "roberta-base"  # fallback
+    if "deberta" in model_name.lower():
+        return DebertaTokenizer.from_pretrained(model_name, cache_dir=TOKENIZER_PATH)
     elif "roberta" in model_name.lower():
-        return RobertaTokenizer.from_pretrained(model_name)
-    elif "deberta" in model_name.lower():
-        return DebertaTokenizer.from_pretrained(model_name)
+        return RobertaTokenizer.from_pretrained(model_name, cache_dir=TOKENIZER_PATH)
+    elif "bert" in model_name.lower():
+        return BertTokenizer.from_pretrained(model_name, cache_dir=TOKENIZER_PATH)
     else:
         raise ValueError(f"Unsupported tokenizer for model: {model_name}")