debug disk space error

interfaces/illframes.py

@@ -21,6 +21,27 @@ domains = {
     "Migration": "migration"
 }
 
+
+# --- DEBUG ---
+import shutil
+
+def convert_size(size):
+    for unit in ['B', 'KB', 'MB', 'GB', 'TB', 'PB']:
+        if size < 1024:
+            return f"{size:.2f} {unit}"
+        size /= 1024
+
+def get_disk_space(path="/"):
+    total, used, free = shutil.disk_usage(path)
+    
+    return {
+        "Total": convert_size(total),
+        "Used": convert_size(used),
+        "Free": convert_size(free)
+    }
+
+# ---
+
 def check_huggingface_path(checkpoint_path: str):
     try:
         hf_api = HfApi(token=HF_TOKEN)
@@ -34,7 +55,24 @@ def build_huggingface_path(domain: str):
 
 def predict(text, model_id, tokenizer_id, label_names):
     device = torch.device("cpu")
-    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN)
+
+    # --- DEBUG ---
+    
+    disk_space = get_disk_space()
+    print("Disk Space Info:")
+    for key, value in disk_space.items():
+        print(f"{key}: {value}")
+        
+    # ---
+
+    try:
+        model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN)
+    except:
+            disk_space = get_disk_space()
+        print("Disk Space Error:")
+        for key, value in disk_space.items():
+            print(f"{key}: {value}")
+        
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,