Merge pull request #3 from shneeba/unloading-model

unload model
llmware-ai · Feb 1, 2025 · 9333e06 · 9333e06
2 parents b5ae53f + 241d456
commit 9333e06
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/llmware/models.py b/llmware/models.py
@@ -10844,8 +10844,12 @@ def load_model(cls, model_name, sample=False, temperature=0.0, get_logits=True,
 
     @classmethod
     def unload_model(cls, model_name):
-        """ Not implemented currently. """
-        return 0
+        if model_name in cls._ModelState.models_list:
+            # Remove the reference so Python can GC it
+            delattr(cls._ModelState, model_name)
+            cls._ModelState.models_list.remove(model_name)
+            cls._ModelState.models_loaded -= 1
+            logger.info(f"ModelResources - Unloaded model: {model_name}")
 
     @classmethod
     def check_if_model_loaded(cls, model_name):