feat: improve error handling for module imports in PretrainedModelForQA

JohnSnowLabs · Jan 27, 2025 · e4f8d7e · e4f8d7e
1 parent d82daca
commit e4f8d7e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/langtest/modelhandler/llm_modelhandler.py b/langtest/modelhandler/llm_modelhandler.py
@@ -165,6 +165,12 @@ def load_model(cls, hub: str, path: str, *args, **kwargs) -> "PretrainedModelFor
                 cls.model = model(**path)
             return cls(hub, cls.model, *args, **filtered_kwargs)
 
+        except ModuleNotFoundError:
+            module = CHAT_MODEL_CLASSES[hub].get("module", None)
+            if module:
+                package_name = module.split(".")[0]
+                raise ValueError(Errors.E078(hub=hub, lib=package_name))
+            raise ValueError(Errors.E078(hub=hub, lib=package_name))
         except ImportError:
             raise ValueError(Errors.E044(path=path))
         except ValidationError as e: