fix: Fix error message when late interaction model name not found (#365)

hh-space-invader · web-flow · commit 58b5a8ed9a13 · 2024-10-14T23:38:52.000+03:00
chore: Updated type hint
diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py
@@ -185,7 +185,7 @@ def embed(
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, List[str]], **kwargs) -> np.ndarray:
+    def query_embed(self, query: Union[str, List[str]], **kwargs) -> Iterable[np.ndarray]:
         if isinstance(query, str):
             query = [query]
 
diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py
@@ -64,8 +64,8 @@ def __init__(
                 return
 
         raise ValueError(
-            f"Model {model_name} is not supported in SparseTextEmbedding."
-            "Please check the supported models using `SparseTextEmbedding.list_supported_models()`"
+            f"Model {model_name} is not supported in LateInteractionTextEmbedding."
+            "Please check the supported models using `LateInteractionTextEmbedding.list_supported_models()`"
         )
 
     def embed(

Original file line number	Diff line number	Diff line change
`@@ -185,7 +185,7 @@ def embed(`
`185`	`185`	`**kwargs,`
`186`	`186`	`)`
`187`	`187`
`188`		`- def query_embed(self, query: Union[str, List[str]], **kwargs) -> np.ndarray:`
	`188`	`+ def query_embed(self, query: Union[str, List[str]], **kwargs) -> Iterable[np.ndarray]:`
`189`	`189`	`if isinstance(query, str):`
`190`	`190`	`query = [query]`
`191`	`191`
Original file line number	Diff line number	Diff line change
`@@ -64,8 +64,8 @@ def __init__(`
`64`	`64`	`return`
`65`	`65`
`66`	`66`	`raise ValueError(`
`67`		`- f"Model {model_name} is not supported in SparseTextEmbedding."`
`68`		- "Please check the supported models using `SparseTextEmbedding.list_supported_models()`"
	`67`	`+ f"Model {model_name} is not supported in LateInteractionTextEmbedding."`
	`68`	+ "Please check the supported models using `LateInteractionTextEmbedding.list_supported_models()`"
`69`	`69`	`)`
`70`	`70`
`71`	`71`	`def embed(`