Refactored model loading in BaseTask class

langtest/tasks/task.py

@@ -50,22 +50,21 @@ def load_model(cls, model_path: str, model_hub: str, *args, **kwargs):
         if "user_prompt" in kwargs:
             cls.user_prompt = kwargs.get("user_prompt")
             kwargs.pop("user_prompt")
-
-        if cls._name in models[model_hub] or cls._name in models["llm"]:
+        try:
+            if model_hub in LANGCHAIN_HUBS:
+                # LLM models
+                cls.model = models["llm"][cls._name].load_model(
+                    hub=model_hub, path=model_path, *args, **kwargs
+                )
+            else:
+                # JSL, Huggingface, and Spacy models
+                cls.model = models[model_hub][cls._name].load_model(
+                    path=model_path, *args, **kwargs
+                )
+            return cls.model
+        except TypeError:
             raise ValueError(Errors.E081.format(hub=model_hub))
 
-        if model_hub in LANGCHAIN_HUBS:
-            # LLM models
-            cls.model = models["llm"][cls._name].load_model(
-                hub=model_hub, path=model_path, *args, **kwargs
-            )
-        else:
-            # JSL, Huggingface, and Spacy models
-            cls.model = models[model_hub][cls._name].load_model(
-                path=model_path, *args, **kwargs
-            )
-        return cls.model
-
     @classmethod
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)