make fastembed cpu or gpu optional (#19878)

Author: logan-markewich

llama-index-integrations/embeddings/llama-index-embeddings-fastembed/llama_index/embeddings/fastembed/base.py

@@ -1,10 +1,12 @@
 import asyncio
 import numpy as np
-from typing import Any, List, Literal, Optional
+from typing import Any, List, Literal, Optional, TYPE_CHECKING
 
 from llama_index.core.base.embeddings.base import BaseEmbedding
 from llama_index.core.bridge.pydantic import Field, PrivateAttr, ConfigDict
-from fastembed import TextEmbedding
+
+if TYPE_CHECKING:
+    from fastembed import TextEmbedding
 
 
 class FastEmbedEmbedding(BaseEmbedding):
@@ -57,7 +59,7 @@ class FastEmbedEmbedding(BaseEmbedding):
         default=None, description="The ONNX providers to use for the embedding model."
     )
 
-    _model: TextEmbedding = PrivateAttr()
+    _model: "TextEmbedding" = PrivateAttr()
 
     def __init__(
         self,
@@ -77,6 +79,15 @@ def __init__(
             **kwargs,
         )
 
+        try:
+            from fastembed import TextEmbedding
+        except ImportError as e:
+            raise ImportError(
+                "Could not import FastEmbed. "
+                "Please install it with `pip install fastembed` or "
+                "`pip install fastembed-gpu` for GPU support"
+            ) from e
+
         self._model = TextEmbedding(
             model_name=model_name,
             cache_dir=cache_dir,

llama-index-integrations/embeddings/llama-index-embeddings-fastembed/pyproject.toml

@@ -4,6 +4,7 @@ build-backend = "hatchling.build"
 
 [dependency-groups]
 dev = [
+    "fastembed>=0.2.2",
     "ipython==8.10.0",
     "jupyter>=1.0.0,<2",
     "mypy==0.991",
@@ -26,14 +27,13 @@ dev = [
 
 [project]
 name = "llama-index-embeddings-fastembed"
-version = "0.4.1"
+version = "0.5.0"
 description = "llama-index embeddings fastembed integration"
 authors = [{name = "Your Name", email = "[email protected]"}]
 requires-python = ">=3.9,<3.13"
 readme = "README.md"
 license = "MIT"
 dependencies = [
-    "fastembed>=0.2.2",
     "llama-index-core>=0.13.0,<0.15",
 ]
 

llama-index-integrations/embeddings/llama-index-embeddings-fastembed/tests/test_embeddings_fastembed.py

@@ -9,7 +9,7 @@ def test_class():
     assert BaseEmbedding.__name__ in names_of_base_classes
 
 
-@patch("llama_index.embeddings.fastembed.base.TextEmbedding")
+@patch("fastembed.TextEmbedding")
 def test_create_fastembed_embedding(mock_text_embedding):
     cache = Path("./test_cache_2")