files

BERTmodel.py

@@ -10,7 +10,7 @@ from transformers import (
     BertForSequenceClassification,
 )
 import torch.nn.functional as F
-
+import os
 
 CACHE_DIR  = "./hf_cache"
  
@@ -18,22 +18,67 @@ CACHE_DIR  = "./hf_cache"
 class BERTVisualizer(TransformerVisualizer):
     def __init__(self,task):
         super().__init__()  
-        print(task,'BERTVIS START')
         self.task = task
-        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', cache_dir=CACHE_DIR)
+        print(task,'BERTVIS START')
+
+        TOKENIZER = 'bert-base-uncased'
+        LOCAL_PATH = os.path.join(CACHE_DIR, "tokenizers",TOKENIZER.replace("/", "_"))
+        
+
+
+        try:
+            self.tokenizer = BertTokenizer.from_pretrained(LOCAL_PATH, local_files_only=True)
+        except Exception as e:
+            self.tokenizer = BertTokenizer.from_pretrained(TOKENIZER)
+            self.tokenizer.save_pretrained(LOCAL_PATH)
+
+
+
         print('finding model', self.task)
         if self.task == 'mlm':
-            self.model = BertForMaskedLM.from_pretrained(
-                "bert-base-uncased",
-                attn_implementation="eager",  # fallback to standard attention
-                cache_dir=CACHE_DIR
-            ).to(self.device)
+            
+            MODEL = 'bert-base-uncased'
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+            
+            try:
+                self.model = BertForMaskedLM.from_pretrained(  LOCAL_PATH, local_files_only=True,   attn_implementation="eager" ).to(self.device)
+            except Exception as e:
+                self.model = BertForMaskedLM.from_pretrained(  MODEL,    attn_implementation="eager" ).to(self.device)
+                self.model.save_pretrained(LOCAL_PATH)
+
         elif self.task == 'sst':
-            self.model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-SST-2",device_map=None, cache_dir=CACHE_DIR)
+            MODEL = "textattack/bert-base-uncased-SST-2"
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+
+            
+            try:
+                self.model = BertForSequenceClassification.from_pretrained(  LOCAL_PATH, local_files_only=True,  device_map=None )
+            except Exception as e:
+                self.model = BertForSequenceClassification.from_pretrained(  MODEL,    device_map=None )
+                self.model.save_pretrained(LOCAL_PATH)
+
+
         elif self.task == 'mnli':
-            self.model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-MNLI", device_map=None, cache_dir=CACHE_DIR)
+            MODEL = 'textattack/bert-base-uncased-MNLI'
+
+            
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+
+            
+            try:
+                self.model = BertForSequenceClassification.from_pretrained(  LOCAL_PATH, local_files_only=True,  device_map=None )
+            except Exception as e:
+                self.model = BertForSequenceClassification.from_pretrained(  MODEL,    device_map=None)
+                self.model.save_pretrained(LOCAL_PATH)
+
+
+ 
+
         else:
             raise ValueError(f"Unsupported task: {self.task}")
+        
+
+        
         print('model found')
         #self.model.to(self.device)
         print('self device junk')

DISTILLBERTmodel.py

@@ -4,7 +4,6 @@ import torch.nn.functional as F
  
 
 import os
-from transformers import DistilBertModel, DistilBertTokenizer
 from models import TransformerVisualizer
 
 from transformers import (
@@ -17,17 +16,65 @@ class DistilBERTVisualizer(TransformerVisualizer):
     def __init__(self, task):
         super().__init__()
         self.task = task
-        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased', cache_dir=CACHE_DIR)
+
+        
+        TOKENIZER = 'distilbert-base-uncased'
+        LOCAL_PATH = os.path.join(CACHE_DIR, "tokenizers",TOKENIZER.replace("/", "_"))
+        
+        try:
+            self.tokenizer = DistilBertTokenizer.from_pretrained(LOCAL_PATH, local_files_only=True)
+        except Exception as e:
+            self.tokenizer = DistilBertTokenizer.from_pretrained(TOKENIZER)
+            self.tokenizer.save_pretrained(LOCAL_PATH)
+
+
+
+        print('finding model', self.task)
         if self.task == 'mlm':
-            self.model = DistilBertForMaskedLM.from_pretrained('distilbert-base-uncased', cache_dir=CACHE_DIR)
+            
+            MODEL = 'distilbert-base-uncased'
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+            
+            try:
+                self.model = DistilBertForMaskedLM.from_pretrained(  LOCAL_PATH, local_files_only=True )
+            except Exception as e:
+                self.model = DistilBertForMaskedLM.from_pretrained(  MODEL  )
+                self.model.save_pretrained(LOCAL_PATH)
+
         elif self.task == 'sst':
-            self.model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased-finetuned-sst-2-english', cache_dir=CACHE_DIR)
+            MODEL = 'distilbert-base-uncased-finetuned-sst-2-english'
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+
+            
+            try:
+                self.model = DistilBertForSequenceClassification.from_pretrained(  LOCAL_PATH, local_files_only=True )
+            except Exception as e:
+                self.model = DistilBertForSequenceClassification.from_pretrained(  MODEL )
+                self.model.save_pretrained(LOCAL_PATH)
+
+
         elif self.task == 'mnli':
-            self.model = DistilBertForSequenceClassification.from_pretrained("textattack/distilbert-base-uncased-MNLI", cache_dir=CACHE_DIR)
+            MODEL = "textattack/distilbert-base-uncased-MNLI"
+
+            
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+
+            
+            try:
+                self.model = DistilBertForSequenceClassification.from_pretrained(  LOCAL_PATH, local_files_only=True)
+            except Exception as e:
+                self.model = DistilBertForSequenceClassification.from_pretrained(  MODEL)
+                self.model.save_pretrained(LOCAL_PATH)
+
 
+ 
 
         else:
-            raise NotImplementedError("Task not supported for DistilBERT")
+            raise ValueError(f"Unsupported task: {self.task}")
+        
+
+  
+ 
         
 
         self.model.eval()

ROBERTAmodel.py

@@ -5,19 +5,60 @@ from models import TransformerVisualizer
 from transformers import (
     RobertaForMaskedLM, RobertaForSequenceClassification
 )
-
+import os
 CACHE_DIR  = "./hf_cache"
 class RoBERTaVisualizer(TransformerVisualizer):
     def __init__(self, task):
         super().__init__()
         self.task = task
-        self.tokenizer = RobertaTokenizer.from_pretrained('roberta-base', cache_dir=CACHE_DIR)
+
+
+        
+        TOKENIZER = 'roberta-base'
+        LOCAL_PATH = os.path.join(CACHE_DIR, "tokenizers",TOKENIZER.replace("/", "_"))
+        
+        try:
+            self.tokenizer = RobertaTokenizer.from_pretrained(LOCAL_PATH, local_files_only=True)
+        except Exception as e:
+            self.tokenizer = RobertaTokenizer.from_pretrained(TOKENIZER)
+            self.tokenizer.save_pretrained(LOCAL_PATH)
+
         if self.task == 'mlm':
-            self.model = RobertaForMaskedLM.from_pretrained("roberta-base", cache_dir=CACHE_DIR)
+            
+            MODEL = "roberta-base"
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+            
+            try:
+                self.model = RobertaForMaskedLM.from_pretrained(  LOCAL_PATH, local_files_only=True )
+            except Exception as e:
+                self.model = RobertaForMaskedLM.from_pretrained(  MODEL  )
+                self.model.save_pretrained(LOCAL_PATH)
+ 
         elif self.task == 'sst':
-            self.model = RobertaForSequenceClassification.from_pretrained('textattack/roberta-base-SST-2', cache_dir=CACHE_DIR)
+
+            
+            MODEL = 'textattack/roberta-base-SST-2'
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+            
+            try:
+                self.model = RobertaForSequenceClassification.from_pretrained(  LOCAL_PATH, local_files_only=True )
+            except Exception as e:
+                self.model = RobertaForSequenceClassification.from_pretrained(  MODEL )
+                self.model.save_pretrained(LOCAL_PATH)
+
+
         elif self.task == 'mnli':
-            self.model = RobertaForSequenceClassification.from_pretrained("roberta-large-mnli", cache_dir=CACHE_DIR)
+            MODEL = "roberta-large-mnli"
+            LOCAL_PATH = os.path.join(CACHE_DIR, "models",MODEL.replace("/", "_"))
+
+            
+            try:
+                self.model = RobertaForSequenceClassification.from_pretrained(  LOCAL_PATH, local_files_only=True)
+            except Exception as e:
+                self.model = RobertaForSequenceClassification.from_pretrained(  MODEL)
+                self.model.save_pretrained(LOCAL_PATH)
+
+
 
 
         self.model.to(self.device)

server.py

@@ -9,18 +9,6 @@ from ROBERTAmodel import *
 from BERTmodel import *
 from DISTILLBERTmodel import *
 
-import shutil
-import os
-
-CACHE_DIR = "./hf_cache"
-if os.path.exists(CACHE_DIR):
-    try:
-        shutil.rmtree(CACHE_DIR)
-        print("✅ Cleared hf_cache directory")
-    except Exception as e:
-        print("❌ Failed to clear hf_cache:", e)
-
-        
 VISUALIZER_CLASSES = {
     "BERT": BERTVisualizer,
     "RoBERTa": RoBERTaVisualizer,
@@ -229,3 +217,19 @@ def get_grad_attn_matrix(req: GradAttnModelRequest):
         return {"error": str(e)}
     
 
+if __name__ == "__main__":
+     
+    print('rim ')
+    BERTVisualizer('mlm')
+    BERTVisualizer('mnli')
+    BERTVisualizer('sst')
+
+
+    RoBERTaVisualizer('mlm')
+    RoBERTaVisualizer('mnli')
+    RoBERTaVisualizer('sst')
+
+
+    DistilBERTVisualizer('mlm')
+    DistilBERTVisualizer('mnli')
+    DistilBERTVisualizer('sst')