Fix HF registry return type.

explosion · Mar 21, 2024 · 91a1ee0 · 91a1ee0
1 parent a97bbe1
commit 91a1ee0
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/spacy_llm/models/hf/registry.py b/spacy_llm/models/hf/registry.py
@@ -1,8 +1,9 @@
-from typing import Any, Callable, Dict, Iterable, Optional
+from typing import Any, Dict, Optional
 
 from confection import SimpleFrozenDict
 
 from ...registry import registry
+from .base import HuggingFace
 from .dolly import Dolly
 from .falcon import Falcon
 from .llama2 import Llama2
@@ -17,7 +18,7 @@ def huggingface_v1(
     name: str,
     config_init: Optional[Dict[str, Any]] = SimpleFrozenDict(),
     config_run: Optional[Dict[str, Any]] = SimpleFrozenDict(),
-) -> Callable[[Iterable[Iterable[str]]], Iterable[Iterable[str]]]:
+) -> HuggingFace:
     """Returns HuggingFace model instance.
     name (str): Name of model to use.
     config_init (Optional[Dict[str, Any]]): HF config for initializing the model.