Custom RAG Workflow
ReadtheDocs Retrieval Augmented Generation (RAG) using Zilliz Free Tier
In this notebook, we are going to use Milvus documentation pages to create a chatbot about our product. The chatbot is going to follow RAG steps to retrieve chunks of data using Semantic Vector Search, then the Question + Context will be fed as a Prompt to a LLM to generate an answer.
Many RAG demos use OpenAI for the Embedding Model and ChatGPT for the Generative AI model. In this notebook, we will demo a fully open source RAG stack.
Using open-source Q&A with retrieval saves money since we make free calls to our own data almost all the time - retrieval, evaluation, and development iterations. We only make a paid call to OpenAI once for the final chat generation step.
Let's get started!
Start up a Zilliz free tier cluster.
Code in this notebook uses fully-managed Milvus on Ziliz Cloud free trial.
- Choose the default "Starter" option when you provision > Create collection > Give it a name > Create cluster and collection.
- On the Cluster main page, copy your
API Keyand store it locally in a .env variable. See note below how to do that. - Also on the Cluster main page, copy the
Public Endpoint URI.
💡 Note: To keep your tokens private, best practice is to use an env variable. See how to save api key in env variable.
In Jupyter, you also need a .env file (in same dir as notebooks) containing lines like this:
- VARIABLE_NAME=value
Type of server: Zilliz Cloud Vector Database(Compatible with Milvus 2.3)
Load the Embedding Model checkpoint and use it to create vector embeddings
Embedding model: We will use the open-source sentence transformers available on HuggingFace to encode the documentation text. We will download the model from HuggingFace and run it locally.
Two model parameters of note below:
- EMBEDDING_DIM refers to the dimensionality or length of the embedding vector. In this case, the embeddings generated for EACH token in the input text will have the SAME length = 1024. This size of embedding is often associated with BERT-based models, where the embeddings are used for downstream tasks such as classification, question answering, or text generation.
- MAX_SEQ_LENGTH is the maximum length the encoder model can handle for input sequences. In this case, if sequences longer than 512 tokens are given to the model, everything longer will be (silently!) chopped off. This is the reason why a chunking strategy is needed to segment input texts into chunks with lengths that will fit in the model's input.
device: cpu
No sentence-transformers model found with name /Users/christybergman/.cache/torch/sentence_transformers/WhereIsAI_UAE-Large-V1. Creating a new one with MEAN pooling.
<class 'sentence_transformers.SentenceTransformer.SentenceTransformer'>
SentenceTransformer(
(0): Transformer({'max_seq_length': 512, 'do_lower_case': False}) with Transformer model: BertModel
(1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
)
model_name: WhereIsAI/UAE-Large-V1
EMBEDDING_DIM: 1024
MAX_SEQ_LENGTH: 512
Create a Milvus collection
You can think of a collection in Milvus like a "table" in SQL databases. The collection will contain the
- Schema (or no-schema Milvus client).
💡 You'll need the vectorEMBEDDING_DIMparameter from your embedding model. Typical values are:- 1024 for sbert embedding models
- 1536 for ada-002 OpenAI embedding models
- Vector index for efficient vector search
- Vector distance metric for measuring nearest neighbor vectors
- Consistency level
In Milvus, transactional consistency is possible; however, according to the CAP theorem, some latency must be sacrificed. 💡 Searching movie reviews is not mission-critical, so
eventuallyconsistent is fine here.
Add a Vector Index
The vector index determines the vector search algorithm used to find the closest vectors in your data to the query a user submits.
Most vector indexes use different sets of parameters depending on whether the database is:
- inserting vectors (creation mode) - vs -
- searching vectors (search mode)
Scroll down the docs page to see a table listing different vector indexes available on Milvus. For example:
- FLAT - deterministic exhaustive search
- IVF_FLAT or IVF_SQ8 - Hash index (stochastic approximate search)
- HNSW - Graph index (stochastic approximate search)
- AUTOINDEX - Automatically determined based on OSS vs Zilliz cloud, type of GPU, size of data.
Besides a search algorithm, we also need to specify a distance metric, that is, a definition of what is considered "close" in vector space. In the cell below, the HNSW search index is chosen. Its possible distance metrics are one of:
- L2 - L2-norm
- IP - Dot-product
- COSINE - Angular distance
💡 Most use cases work better with normalized embeddings, in which case L2 is useless (every vector has length=1) and IP and COSINE are the same. Only choose L2 if you plan to keep your embeddings unnormalized.
Successfully dropped collection: `wikipedia` Successfully created collection: `wikipedia`
Insert data into Milvus
For each original text chunk, we'll write the quadruplet (vector, text, source, h1, h2) into the database.
The Milvus Client wrapper can only handle loading data from a list of dictionaries.
Otherwise, in general, Milvus supports loading data from:
- pandas dataframes
- list of dictionaries
Below, we use the embedding model provided by HuggingFace, download its checkpoint, and run it locally as the encoder.
Num docs: 1 Num chunks: 704 Start inserting entities
100%|██████████| 1/1 [00:03<00:00, 3.95s/it]
Milvus Client insert time for 704 vectors: 3.9572505950927734 seconds
Define Evaluation Metrics
/Users/christybergman/mambaforge/envs/py311new/lib/python3.11/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The class `langchain_community.chat_models.openai.ChatOpenAI` was deprecated in langchain-community 0.0.10 and will be removed in 0.2.0. An updated version of the class exists in the langchain-openai package and should be used instead. To use it run `pip install -U langchain-openai` and import as `from langchain_openai import ChatOpenAI`. warn_deprecated(
Define a Custom Execution Loop for RAG.
question = How did New York City get its name? STEP 2: Retrieval from collection #1 MilvusDocs. DISTANCE SCORE: 0.39108937978744507 branching logic... STEP 3: Score is too low, GET INTENT from the user's question. intent = new_york STEP 4: Based on question intent, retrieve from collection #2 Wikipedia. chunk_answer: New York City traces its origins to Fort Amsterdam and a trading post founded on the southern tip of Manhattan Island by Dutch colonists in approximat DISTANCE SCORE: 0.7961502075195312 branch logic... Score from custom RAG Retrieval is above threshold, proceed to answer generation step. STEP 5: Generating GPT3.5 answer from the custom execution loop for RAG in the ASSISTANT PROMPT.
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks... To disable this warning, you can either: - Avoid using `tokenizers` before the fork if possible - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
STEP 6: Evaluate whether the chatbot response answers the initial user query well. evaluating with [answer_similarity]
100%|██████████| 1/1 [00:00<00:00, 1.49it/s]
evaluating with [answer_relevancy]
100%|██████████| 1/1 [00:01<00:00, 1.73s/it]
evaluating with [answer_correctness]
100%|██████████| 1/1 [00:05<00:00, 5.98s/it]
Ragas evaluation: answer similarity: 0.9421961714808575, answer relevancy: 0.894, answer correctness: 0.664
STEP 7: LLM answer passed Evaluation, return it to the user.
('Answer: New York City was originally named New Amsterdam by Dutch colonists '
'in 1626. However, it was renamed New York in 1664 after King Charles II '
'granted the lands to his brother, the Duke of York, when the city came under '
'British control.')
Final Eval Comparisons Custom RAG vs OpenAI RAG
evaluating with [context_recall]
100%|██████████| 1/1 [00:14<00:00, 14.62s/it]
evaluating with [context_precision]
100%|██████████| 1/1 [00:07<00:00, 7.86s/it]
evaluating with [faithfulness]
100%|██████████| 1/1 [00:29<00:00, 29.35s/it]
evaluating with [answer_similarity]
100%|██████████| 1/1 [00:01<00:00, 1.20s/it]
evaluating with [answer_relevancy]
100%|██████████| 1/1 [00:07<00:00, 7.96s/it]
evaluating with [answer_correctness]
100%|██████████| 1/1 [00:20<00:00, 20.12s/it]
evaluating with [answer_similarity]
100%|██████████| 1/1 [00:00<00:00, 2.01it/s]
evaluating with [answer_relevancy]
100%|██████████| 1/1 [00:07<00:00, 7.85s/it]
evaluating with [answer_correctness]
100%|██████████| 1/1 [00:14<00:00, 14.49s/it]
####### FINAL SCORES OPENAI RAG vs MILVUS CUSTOM RAG ######### LLM as judge model: gpt-3.5-turbo-1106 with temperature: 0.1 scores: # Truth vs RAG answers: 4 avg_similarity_Custom_RAG: 0.83 avg_similarity_OpenAI_RAG: 0.78 answer_relevancy_Custom_RAG: 0.96 avg_relevancy_OpenAI_RAG: 0.73 avg_correctness_Custom_RAG: 0.6 avg_correctness_OpenAI_RAG: 0.32
Author: Christy Bergman Python implementation: CPython Python version : 3.11.6 IPython version : 8.18.1 torch : 2.1.1 transformers : 4.35.2 sentence_transformers: 2.2.2 pymilvus : 2.3.4 langchain : 0.1.0 openai : 1.7.2 conda environment: py311new