debug: disk space

interfaces/ontolisst.py

@@ -15,6 +15,26 @@ languages = [
 
 from label_dicts import ONTOLISST_LABEL_NAMES
 
+# --- 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 build_huggingface_path(language: str):
     return "poltextlab/xlm-roberta-large_ontolisst_v1"
@@ -23,6 +43,16 @@ def predict(text, model_id, tokenizer_id):
     device = torch.device("cpu")
     model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+
+    # --- DEBUG ---
+    
+    disk_space = get_disk_space('/data/')
+    print("Disk Space Info:")
+    for key, value in disk_space.items():
+        print(f"{key}: {value}")
+        
+    # ---
+    
     model.to(device)
 
     inputs = tokenizer(text,