jit pt2

utils.py

@@ -57,7 +57,6 @@ for domain in domains_illframes.values():
 tokenizers = ["xlm-roberta-large"]
 
 def download_hf_models():
-    return ### DEBUG
     # Ensure the JIT model directory exists
     os.makedirs(JIT_DIR, exist_ok=True)
 
@@ -107,14 +106,30 @@ def df_h():
     print(result.stdout)
     
 def scan_cache():
+    # Scan Hugging Face model cache
     cache_dir = os.environ.get("TRANSFORMERS_CACHE", os.path.expanduser("~/.cache/huggingface/transformers"))
     scan_result = scan_cache_dir(cache_dir)
     
-    print("=== Model Cache Report ===")
+    print("=== 🤗 Hugging Face Model Cache ===")
     print(f"Cache size: {scan_result.size_on_disk / 1e6:.2f} MB")
     print(f"Number of repos: {len(scan_result.repos)}")
     for repo in scan_result.repos:
         print(f"- {repo.repo_id} ({repo.repo_type}) — {repo.size_on_disk / 1e6:.2f} MB")
+
+    print("\n=== 🧊 TorchScript JIT Cache ===")
+    if not os.path.exists(JIT_DIR):
+        print(f"(Directory does not exist: {JIT_DIR})")
+        return
+
+    total_size = 0
+    for filename in os.listdir(JIT_DIR):
+        if filename.endswith(".pt"):
+            path = os.path.join(JIT_DIR, filename)
+            size = os.path.getsize(path)
+            total_size += size
+            print(f"- {filename}: {size / 1e6:.2f} MB")
+    
+    print(f"Total JIT cache size: {total_size / 1e6:.2f} MB")
     
 def set_hf_cache_dir(path:str):
     os.environ['TRANSFORMERS_CACHE'] = path