dcachedir

BERTmodel.py

@@ -11,23 +11,26 @@ from transformers import (
 )
 import torch.nn.functional as F
 
+
+CACHE_DIR  = "./hf_cache"
  
     
 class BERTVisualizer(TransformerVisualizer):
     def __init__(self,task):
         super().__init__()  
         self.task = task
-        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', cache_dir=CACHE_DIR)
         print('finding model', self.task)
         if self.task == 'mlm':
             self.model = BertForMaskedLM.from_pretrained(
                 "bert-base-uncased",
-                attn_implementation="eager"  # fallback to standard attention
+                attn_implementation="eager",  # fallback to standard attention
+                cache_dir=CACHE_DIR
             ).to(self.device)
         elif self.task == 'sst':
-            self.model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-SST-2",device_map=None)
+            self.model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-SST-2",device_map=None, cache_dir=CACHE_DIR)
         elif self.task == 'mnli':
-            self.model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-MNLI", device_map=None)
+            self.model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-MNLI", device_map=None, cache_dir=CACHE_DIR)
         else:
             raise ValueError(f"Unsupported task: {self.task}")
         print('model found')

DISTILLBERTmodel.py

@@ -12,17 +12,18 @@ from transformers import (
     DistilBertForMaskedLM, DistilBertForSequenceClassification
 )    
 
+CACHE_DIR  = "./hf_cache"
 class DistilBERTVisualizer(TransformerVisualizer):
     def __init__(self, task):
         super().__init__()
         self.task = task
-        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
+        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased', cache_dir=CACHE_DIR)
         if self.task == 'mlm':
-            self.model = DistilBertForMaskedLM.from_pretrained('distilbert-base-uncased')
+            self.model = DistilBertForMaskedLM.from_pretrained('distilbert-base-uncased', cache_dir=CACHE_DIR)
         elif self.task == 'sst':
-            self.model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased-finetuned-sst-2-english')
+            self.model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased-finetuned-sst-2-english', cache_dir=CACHE_DIR)
         elif self.task == 'mnli':
-            self.model = DistilBertForSequenceClassification.from_pretrained("textattack/distilbert-base-uncased-MNLI")
+            self.model = DistilBertForSequenceClassification.from_pretrained("textattack/distilbert-base-uncased-MNLI", cache_dir=CACHE_DIR)
 
 
         else:

ROBERTAmodel.py

@@ -6,17 +6,18 @@ from transformers import (
     RobertaForMaskedLM, RobertaForSequenceClassification
 )
 
+CACHE_DIR  = "./hf_cache"
 class RoBERTaVisualizer(TransformerVisualizer):
     def __init__(self, task):
         super().__init__()
         self.task = task
-        self.tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+        self.tokenizer = RobertaTokenizer.from_pretrained('roberta-base', cache_dir=CACHE_DIR)
         if self.task == 'mlm':
-            self.model = RobertaForMaskedLM.from_pretrained("roberta-base")
+            self.model = RobertaForMaskedLM.from_pretrained("roberta-base", cache_dir=CACHE_DIR)
         elif self.task == 'sst':
-            self.model = RobertaForSequenceClassification.from_pretrained('textattack/roberta-base-SST-2')
+            self.model = RobertaForSequenceClassification.from_pretrained('textattack/roberta-base-SST-2', cache_dir=CACHE_DIR)
         elif self.task == 'mnli':
-            self.model = RobertaForSequenceClassification.from_pretrained("roberta-large-mnli")
+            self.model = RobertaForSequenceClassification.from_pretrained("roberta-large-mnli", cache_dir=CACHE_DIR)
 
 
         self.model.to(self.device)