removed offload disk + added torch memory cleanup

interfaces/cap.py

@@ -11,7 +11,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import CAP_NUM_DICT, CAP_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -86,7 +86,7 @@ def build_huggingface_path(language: str, domain: str):
 #@spaces.GPU
 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).to(device)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN).to(device)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,
@@ -98,6 +98,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     output_pred = {f"[{CAP_NUM_DICT[i]}] {CAP_LABEL_NAMES[CAP_NUM_DICT[i]]}": probs[i] for i in np.argsort(probs)[::-1]}

interfaces/cap_media_demo.py

@@ -10,7 +10,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import CAP_MEDIA_NUM_DICT, CAP_MEDIA_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -35,7 +35,7 @@ def build_huggingface_path(language: str, domain: str):
 
 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)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,
@@ -47,6 +47,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     output_pred = {f"[{CAP_MEDIA_NUM_DICT[i]}] {CAP_MEDIA_LABEL_NAMES[CAP_MEDIA_NUM_DICT[i]]}": probs[i] for i in np.argsort(probs)[::-1]}

interfaces/cap_minor.py

@@ -10,7 +10,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import CAP_MIN_NUM_DICT, CAP_MIN_LABEL_NAMES, CAP_LABEL_NAMES 
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 from itertools import islice
 
 def take(n, iterable):
@@ -67,7 +67,7 @@ def build_huggingface_path(language: str, domain: str):
 
 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)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,
@@ -79,6 +79,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     output_pred = {f"[{'999' if str(CAP_MIN_NUM_DICT[i]) == '999' else str(CAP_MIN_NUM_DICT[i])[:-2]}]{convert_minor_to_major(CAP_MIN_NUM_DICT[i])} [{CAP_MIN_NUM_DICT[i]}]{CAP_MIN_LABEL_NAMES[CAP_MIN_NUM_DICT[i]]}": probs[i] for i in np.argsort(probs)[::-1]}

interfaces/cap_minor_media.py

@@ -15,7 +15,7 @@ from label_dicts import (CAP_MEDIA_NUM_DICT, CAP_MEDIA_LABEL_NAMES,
                         CAP_MIN_NUM_DICT, CAP_MIN_LABEL_NAMES,
                         CAP_MIN_MEDIA_NUM_DICT)
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -116,6 +116,8 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
             minor_logits = minor_model(**inputs).logits
             minor_probs = F.softmax(minor_logits, dim=-1)
             
+        free_gpu_memory(major_model, major_model_id)
+        free_gpu_memory(minor_model, minor_model_id)
         
         print(minor_probs) # debug
         # Restrict to valid minor codes
@@ -148,7 +150,7 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
 
 def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
     device = torch.device("cpu")
-    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN).to(device)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN).to(device)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,
@@ -160,6 +162,7 @@ def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     top_indices = np.argsort(probs)[::-1][:10]

interfaces/emotion.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import MANIFESTO_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -27,7 +27,7 @@ def build_huggingface_path(language: str):
 
 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)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
     model.to(device)
 
@@ -40,6 +40,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     output_pred = {model.config.id2label[i]: probs[i] for i in np.argsort(probs)[::-1]}

interfaces/emotion9.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import EMOTION9_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -38,6 +38,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 

interfaces/illframes.py

@@ -10,7 +10,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import ILLFRAMES_MIGRATION_LABEL_NAMES, ILLFRAMES_COVID_LABEL_NAMES, ILLFRAMES_WAR_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -59,7 +59,7 @@ def build_huggingface_path(domain: str):
 def predict(text, model_id, tokenizer_id, label_names):
     device = torch.device("cpu")
     try:
-        model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, offload_folder="offload", device_map="auto", token=HF_TOKEN)
+        model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     except:
         disk_space = get_disk_space('/data/')
         print("Disk Space Error:")
@@ -80,6 +80,7 @@ def predict(text, model_id, tokenizer_id, label_names):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 

interfaces/manifesto.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import MANIFESTO_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -26,7 +26,7 @@ def build_huggingface_path(language: str):
 
 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)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,
@@ -38,6 +38,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     output_pred = {f"[{model.config.id2label[i]}] {MANIFESTO_LABEL_NAMES[int(model.config.id2label[i])]}": probs[i] for i in np.argsort(probs)[::-1]}

interfaces/ontolisst.py

@@ -15,7 +15,7 @@ languages = [
 
 from label_dicts import ONTOLISST_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 # --- DEBUG ---
 import shutil
@@ -41,9 +41,10 @@ def get_disk_space(path="/"):
 def build_huggingface_path(language: str):
     return "poltextlab/xlm-roberta-large_ontolisst_v1"
 
+
 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)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     # --- DEBUG ---
@@ -66,6 +67,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     predicted_class_id = probs.argmax()
@@ -78,6 +80,7 @@ def predict(text, model_id, tokenizer_id):
     output_info = f'<p style="text-align: center; display: block">Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.</p>'
     return output_pred, output_info
 
+
 def predict_cap(text, language):
     model_id = build_huggingface_path(language)
     tokenizer_id = "xlm-roberta-large"

interfaces/sentiment.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi
 
 from label_dicts import MANIFESTO_LABEL_NAMES
 
-from .utils import is_disk_full
+from .utils import is_disk_full, free_gpu_memory
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -30,7 +30,7 @@ def build_huggingface_path(language: str):
 
 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)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
     model.to(device)
 
@@ -43,6 +43,7 @@ def predict(text, model_id, tokenizer_id):
 
     with torch.no_grad():
         logits = model(**inputs).logits
+    free_gpu_memory(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     predicted_class_id = probs.argmax()

interfaces/utils.py

@@ -1,4 +1,7 @@
 import shutil
+import torch
+import gc
+
 
 def is_disk_full(min_free_space_in_GB=10):
     total, used, free = shutil.disk_usage("/")
@@ -9,4 +12,28 @@ def is_disk_full(min_free_space_in_GB=10):
         return False
     else:
         print('clean up!')
-        return True
+        return True
+        
+        
+def free_gpu_memory(obj=None, label='Object'):
+    """
+    Frees GPU memory by deleting the passed object, collecting garbage,
+    and clearing PyTorch's CUDA cache.
+
+    Args:
+        obj (optional): The model or pipeline to delete.
+        label (str): Label for logging (e.g., 'model', 'pipeline').
+    """
+    if torch.cuda.is_available():
+        print(f"\n[{label}] Before deletion: {torch.cuda.memory_allocated() / 1e6:.2f} MB")
+    
+    if obj is not None:
+        del obj
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    if torch.cuda.is_available():
+        print(f"[{label}] After deletion:  {torch.cuda.memory_allocated() / 1e6:.2f} MB\n")
+        
+

utils.py

@@ -54,8 +54,7 @@ tokenizers = ["xlm-roberta-large"]
 
 def download_hf_models():
     for model_id in models:
-        AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", 
-                                                                   token=HF_TOKEN)
+        AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
     for tokenizer_id in tokenizers:
         AutoTokenizer.from_pretrained(tokenizer_id)