Using PyMilvus's Model To Generate Text Embeddings

Quickly enhance your search capabilities with text embeddings using PyMilvus(higher than 2.4.0). This guide shows how to utilize PyMilvus models for extracting rich text embeddings, setting the foundation for powerful search functionalities.
In this doc, we will go through dense embedding models, sparse embedding models, and hybrid models to show how to use them in action.
Let's begin by installing the dependencies(you could also use virtualenv to create a new enviroment):

For most use cases, you can generate embeddings for storage or retrieval by simply using ef(texts). However, when you need different processing for queries and documents, you could use two specific functions. Documents are processed with encode_documents to generate their embeddings, which are then stored in the vector database. For retrieval, the query is processed using encode_queries to create its embedding, which is then used to search the database.

Dense embedding is a technique used in natural language processing to represent words or phrases as continuous, dense vectors in a high-dimensional space, capturing semantic relationships.

OpenAI Embedding Function

OpenAI offers dense embedding services, but to access them, users must sign up and obtain an API key. With the API key properly set in your environment variables, you can start using tools like OpenAIEmbeddingFunction to generate dense embeddings.

when using OpenAIEmbeddingFunction, encoding_queries and encoding_documents are exactly same procedure, one can use openai_ef(texts) instead.

• openai_ef(texts): same with other two functions.
• openai_ef.encode_queries(queries): same with other two functions.
• openai_ef.encode_documents(documents): same with other two functions.

Additionally, you may initialize the OpenAIEmbeddingFunction by directly providing OpenAI's official parameters, such as api_key and base_url, as part of the function's configuration.

Sentence Transformer Embedding Function

In addition to hosted services like OpenAI, there exists a variety of powerful open-source dense embedding models. For these, the SentenceTransformerEmbeddingFunction based on Sentence-Transformer can be utilized to extract text embeddings.

when using SentenceTransformerEmbeddingFunction, encoding_queries and encoding_documents would prepend query_instruction and doc_instruction respectively, others are same.

• sentence_transformer_ef(texts): prepend nothing, just process the raw text.
• sentence_transformer_ef.encode_queries(queries): prepend the query_instruction to each query.
• sentence_transformer_ef.encode_documents(documents): prepend the doc_instruction to each document.

Additionally, the initialization of SentenceTransformerEmbeddingFunction may incorporate features from Sentence Transformer, such as specifying parameters like the batch_size. Some models require adding an instruction before the actual text input.

Sparse embedding represents words or phrases using vectors where most elements are zero, with only one non-zero element indicating the presence of a specific word in a vocabulary. Sparse embeddings models are efficient and interpretable, making them suitable for tasks where exact term matches are crucial

Splade Embedding Function

SPLADE embedding is a model that offers highly sparse representations for documents and queries, inheriting desirable properties from bag-of-words (BOW) models such as exact term matching and efficiency. We can use SPLADE model easily with SpladeEmbeddingFunction.

when using SpladeEmbeddingFunbction, encoding_queries and encoding_documents would prepend query_instruction and doc_instruction respectively, and k_tokens_query is used to prune the query results. k_tokens_document is used to prune the document results.

• splade_ef(texts): prepend nothing, it does not prune the results.
• splade_ef.encode_queries(queries): prepend the query_instruction to each query, k_tokens_query is used to prune the query results.
• splade_ef.encode_documents(documents): prepend the doc_instruction to each document, k_tokens_document is used to prune the document results.

By default, the model outputs the result directly. However, there are situations where users may only want to retain the top k largest values for a desired valid value. In such cases, users can specify the parameters 'k_tokens_query' and 'k_tokens_document' for queries and documents, respectively.

BM25 Embedding Function

BM25 is a ranking function used in information retrieval to estimate the relevance of documents to a given search query. It enhances the basic term frequency approach by incorporating document length normalization and term frequency saturation. BM25 can generate sparse embeddings by representing documents as vectors of term importance scores, allowing for efficient retrieval and ranking in sparse vector spaces.

The BM25 algorithm processes text by first breaking it into tokens using a built-in analyzer, as shown with English language tokens like 'artifici', 'intellig', and 'academ'. It then gathers statistics on these tokens, evaluating their frequency and distribution across documents. The core of BM25 calculates the relevance score of each token based on its importance, with rarer tokens receiving higher scores. This concise process enables effective ranking of documents by relevance to a query.

So we need to fit a dataset(or corpus) to get the statistics.

when using BM25EmbeddingFunction, encoding_queries and encoding_documents are not exchangable mathematically. There is no avaliable bm25_ef(texts) implmented.

• bm25_ef(texts): is not avaliable.
• bm25_ef.encode_queries(queries): has its distinct implmentation.
• bm25_ef.encode_documents(documents): has its distinct implmentation.

Fitting data each time is time-consuming; we offer save and load features for efficiency.

Calculating statistics within similar distributions is crucial for achieving accurate results. But when there lacks corpus, we provide pre-built, fitted data on the MS MARCO dataset specifically for English. This will download a pre-built json file into your local path.

Finally, let's show analyzers for other languages.

In the realm of embedding models, hybrid architectures exist that generate both dense and sparse embeddings. We refer to these models as hybrid models and introduce BGE-M3 as an example of such a model.

BGE-M3 Embedding Function

The BGE-M3 is named for its capabilities in Multi-Linguality, Multi-Functionality, and Multi-Granularity. Capable of supporting over 100 languages, BGE-M3 sets new benchmarks in multi-lingual and cross-lingual retrieval tasks. Its unique ability to perform dense retrieval, multi-vector retrieval, and sparse retrieval within a single framework makes it an ideal choice for a wide range of information retrieval (IR) applications.

(caution: when running this section, please restart the jupyter python kernel firstly.)

when using BGEM3EmbeddingFunction, encoding_queries and encoding_documents are exactly same procedure, one can use bge_m3_ef(texts) instead.

• bge_m3_ef(texts): same with other two functions.
• bge_m3_ef.encode_queries(queries): same with other two functions.
• bge_m3_ef.encode_documents(documents): same with other two functions.

Although BGE-M3 is capable of generating both dense and sparse embeddings simultaneously, it can also be configured to function as a standard dense or sparse embedding generator by adjusting its return options.

Similarly, we can utilize the BGEM3EmbeddingFunction as a sparse embedding function.