Merge branch 'main' of https://huggingface.co/spaces/yifan0sun/BERTGradGraph

server.py

@@ -9,6 +9,11 @@ from ROBERTAmodel import *
 from BERTmodel import *
 from DISTILLBERTmodel import *
 
+import os
+import zipfile
+import shutil
+
+
 VISUALIZER_CLASSES = {
     "BERT": BERTVisualizer,
     "RoBERTa": RoBERTaVisualizer,
@@ -57,6 +62,105 @@ def ping():
     return {"message": "pong"}
 
 
+
+    
+def zip_if_needed(src_dir, zip_path):
+    if os.path.exists(zip_path):
+        return  # already zipped
+    print(f"📦 Zipping {src_dir} → {zip_path}")
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for root, _, files in os.walk(src_dir):
+            for file in files:
+                full_path = os.path.join(root, file)
+                rel_path = os.path.relpath(full_path, src_dir)
+                zf.write(full_path, arcname=rel_path)
+    print(f"✅ Created zip: {zip_path}")
+
+def extract_zip_if_needed(zip_path, dest_dir):
+    if os.path.exists(dest_dir):
+        print(f"✅ Already exists: {dest_dir}")
+        return
+    print(f"🔓 Extracting {zip_path} → {dest_dir}")
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(dest_dir)
+    print(f"✅ Extracted to: {dest_dir}")
+
+def print_directory_tree(path):
+    printstr = ''
+    for root, dirs, files in os.walk(path):
+        indent = '  ' * (root.count(os.sep) - path.count(os.sep))
+        printstr += indent
+        printstr += os.path.basename(root)
+        printstr += '\n'
+        for f in files:
+            printstr += indent
+            printstr += f
+    return printstr
+
+
+def copy_zip_extract_and_report():
+    src_base = "./hf_cache"
+    dst_base = "/data/hf_cache"
+
+    for category in ["models", "tokenizers"]:
+        src_dir = os.path.join(src_base, category)
+        dst_dir = os.path.join(dst_base, category)
+
+        if not os.path.exists(src_dir):
+            continue
+
+        os.makedirs(dst_dir, exist_ok=True)
+
+        for name in os.listdir(src_dir):
+            full_path = os.path.join(src_dir, name)
+            if not os.path.isdir(full_path):
+                continue
+
+            zip_name = f"{name}.zip"
+            local_zip = os.path.join(src_dir, zip_name)
+            dst_zip = os.path.join(dst_dir, zip_name)
+            extract_path = os.path.join(dst_dir, name)
+
+            zip_if_needed(full_path, local_zip)
+
+            # Copy zip to /data
+            if not os.path.exists(dst_zip):
+                shutil.copy(local_zip, dst_zip)
+                print(f"📤 Copied zip to: {dst_zip}")
+
+            extract_zip_if_needed(dst_zip, extract_path)
+
+    print("\n📦 Local hf_cache structure:")
+    printstr1 = print_directory_tree("./hf_cache")
+
+    print("\n💾 Persistent /data/hf_cache structure:")
+    printstr2 = print_directory_tree("/data/hf_cache")
+    return printstr1 + '\n\n' + printstr2
+
+
+
+@app.get("/copy_and_extract")
+def copy_and_extract():
+    import io
+    from contextlib import redirect_stdout
+
+    printstr = copy_zip_extract_and_report()
+    
+    return {"message": "done", "log": printstr}
+
+
+
+
+
+
+
+
+
+
+
+
+    
+
 @app.post("/load_model")
 def load_model(req: LoadModelRequest):
     print(f"\n--- /load_model request received ---")
@@ -216,20 +320,3 @@ def get_grad_attn_matrix(req: GradAttnModelRequest):
         print("SERVER EXCEPTION:", e)
         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')