llama-index

1
from typing import Any, List, Literal, Optional
2

3
import numpy as np
4

5
from llama_index.legacy.bridge.pydantic import Field, PrivateAttr
6
from llama_index.legacy.embeddings.base import BaseEmbedding
7

8

9
class FastEmbedEmbedding(BaseEmbedding):
10
    """
11
    Qdrant FastEmbedding models.
12
    FastEmbed is a lightweight, fast, Python library built for embedding generation.
13
    See more documentation at:
14
    * https://github.com/qdrant/fastembed/
15
    * https://qdrant.github.io/fastembed/.
16

17
    To use this class, you must install the `fastembed` Python package.
18

19
    `pip install fastembed`
20
    Example:
21
        from llama_index.legacy.embeddings import FastEmbedEmbedding
22
        fastembed = FastEmbedEmbedding()
23
    """
24

25
    model_name: str = Field(
26
        "BAAI/bge-small-en-v1.5",
27
        description="Name of the FastEmbedding model to use.\n"
28
        "Defaults to 'BAAI/bge-small-en-v1.5'.\n"
29
        "Find the list of supported models at "
30
        "https://qdrant.github.io/fastembed/examples/Supported_Models/",
31
    )
32

33
    max_length: int = Field(
34
        512,
35
        description="The maximum number of tokens. Defaults to 512.\n"
36
        "Unknown behavior for values > 512.",
37
    )
38

39
    cache_dir: Optional[str] = Field(
40
        None,
41
        description="The path to the cache directory.\n"
42
        "Defaults to `local_cache` in the parent directory",
43
    )
44

45
    threads: Optional[int] = Field(
46
        None,
47
        description="The number of threads single onnxruntime session can use.\n"
48
        "Defaults to None",
49
    )
50

51
    doc_embed_type: Literal["default", "passage"] = Field(
52
        "default",
53
        description="Type of embedding to use for documents.\n"
54
        "'default': Uses FastEmbed's default embedding method.\n"
55
        "'passage': Prefixes the text with 'passage' before embedding.\n"
56
        "Defaults to 'default'.",
57
    )
58

59
    _model: Any = PrivateAttr()
60

61
    @classmethod
62
    def class_name(self) -> str:
63
        return "FastEmbedEmbedding"
64

65
    def __init__(
66
        self,
67
        model_name: Optional[str] = "BAAI/bge-small-en-v1.5",
68
        max_length: Optional[int] = 512,
69
        cache_dir: Optional[str] = None,
70
        threads: Optional[int] = None,
71
        doc_embed_type: Literal["default", "passage"] = "default",
72
    ):
73
        super().__init__(
74
            model_name=model_name,
75
            max_length=max_length,
76
            threads=threads,
77
            doc_embed_type=doc_embed_type,
78
        )
79
        try:
80
            from fastembed.embedding import FlagEmbedding
81

82
            self._model = FlagEmbedding(
83
                model_name=model_name,
84
                max_length=max_length,
85
                cache_dir=cache_dir,
86
                threads=threads,
87
            )
88
        except ImportError as ie:
89
            raise ImportError(
90
                "Could not import 'fastembed' Python package. "
91
                "Please install it with `pip install fastembed`."
92
            ) from ie
93

94
    def _get_text_embedding(self, text: str) -> List[float]:
95
        embeddings: List[np.ndarray]
96
        if self.doc_embed_type == "passage":
97
            embeddings = list(self._model.passage_embed(text))
98
        else:
99
            embeddings = list(self._model.embed(text))
100
        return embeddings[0].tolist()
101

102
    def _get_query_embedding(self, query: str) -> List[float]:
103
        query_embeddings: np.ndarray = next(self._model.query_embed(query))
104
        return query_embeddings.tolist()
105

106
    async def _aget_query_embedding(self, query: str) -> List[float]:
107
        return self._get_query_embedding(query)
108