Update README.md (#62)
Update retrieval corpus preparation in rag/README.md
This commit is contained in:
@@ -1,5 +1,7 @@
|
||||
# RAG Agent Example
|
||||
|
||||
## Overview
|
||||
|
||||
This example originally runs on a single node with four GPUs, each requiring at least 40GB of memory.
|
||||
|
||||
1. Prepare the RAG dataset in the wiki_retriever_mcp folder. Wiki chunks (`nq_list.pkl`) and Faiss index (`nq_hnsw_faiss_n32e40.index`) are required. (Full wiki dump files are huge, additional information will be provided later)
|
||||
@@ -10,6 +12,103 @@ This example originally runs on a single node with four GPUs, each requiring at
|
||||
6. Run the agent: `python rag_agent.py`. This automatically launches 12 agent workers by default.
|
||||
7. In another terminal, launch the training server: `bash train.sh`.
|
||||
|
||||
|
||||
## Preparing the Retrieval Corpus
|
||||
|
||||
To enable semantic retrieval with this mcp server, we need two files:
|
||||
|
||||
1. **FAISS index file** (`.index`)
|
||||
2. **Chunk list file** (`.pkl`)
|
||||
|
||||
These two files work together: the FAISS index stores the vector embeddings and their mapping to integer IDs, while the pickle file stores the actual text chunks. The integer IDs in the index correspond exactly to the positions in the chunk list.
|
||||
|
||||
---
|
||||
|
||||
### Step 1. Collecting Text Chunks
|
||||
|
||||
You first need a collection of text passages (chunks). For example, you can download a Wikipedia-based dataset such as `wiki18_100w.zip` in the [FlashRAG_dataset](https://huggingface.co/datasets/FlashRAG) or use other pre-split corpora.
|
||||
|
||||
---
|
||||
|
||||
### Step 2. Creating the FAISS Index (`nq_hnsw_faiss_n32e40.index`)
|
||||
|
||||
- Use a sentence embedding model (e.g., `BAAI/bge-large-en-v1.5`) to encode each chunk into a vector.
|
||||
- Build a FAISS index from these vectors.
|
||||
- In this example, we use an **HNSW index** (Hierarchical Navigable Small World graph), which supports efficient approximate nearest-neighbor search.
|
||||
- The index only stores embeddings and integer IDs (no raw text).
|
||||
|
||||
---
|
||||
|
||||
### Step 3. Creating the Chunk List (`nq_list.pkl`)
|
||||
|
||||
- Store the raw text chunks in a Python list.
|
||||
- Save this list with `pickle`.
|
||||
- The index ID returned by FAISS corresponds to the list index in this file. For example, if FAISS search returns `I[0][i] = 12345`, then the corresponding text chunk is `chunks[12345]`.
|
||||
|
||||
---
|
||||
|
||||
### Example Schema
|
||||
|
||||
- **`nq_hnsw_faiss_n32e40.index`**
|
||||
- Type: FAISS HNSW index
|
||||
- Contains:
|
||||
- Vector embeddings
|
||||
- Graph structure for fast search
|
||||
- Integer IDs mapping to chunk positions
|
||||
|
||||
- **`nq_list.pkl`**
|
||||
- Type: Pickled Python list
|
||||
- Element type: string (or dict with text + metadata, depending on preprocessing)
|
||||
- Example:
|
||||
```python
|
||||
[
|
||||
"The Eiffel Tower is located in Paris, France.",
|
||||
"Albert Einstein developed the theory of relativity.",
|
||||
...
|
||||
]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Step 4. Code Example: Building Index and Chunk List
|
||||
Warning: The following example only demonstrates a small-scale workflow. In practice, if the dataset is large, you should encode the text in batches and incrementally add them to the index.
|
||||
|
||||
```python
|
||||
import faiss
|
||||
import pickle
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
# 1. Prepare your text chunks (list of strings)
|
||||
chunk_texts = [
|
||||
"The Eiffel Tower is located in Paris, France.",
|
||||
"Albert Einstein developed the theory of relativity.",
|
||||
"Python is a popular programming language.",
|
||||
# ... more chunks
|
||||
]
|
||||
|
||||
# 2. Load embedding model
|
||||
model = SentenceTransformer("BAAI/bge-large-en-v1.5")
|
||||
|
||||
# 3. Encode text chunks into embeddings
|
||||
embeddings = model.encode(chunk_texts, normalize_embeddings=True)
|
||||
|
||||
# 4. Build FAISS HNSW index
|
||||
dim = embeddings.shape[1]
|
||||
index = faiss.IndexHNSWFlat(dim, 32) # 32 neighbors by default
|
||||
index.hnsw.efConstruction = 40 # efConstruction parameter
|
||||
index.add(embeddings)
|
||||
|
||||
# 5. Save FAISS index
|
||||
faiss.write_index(index, "nq_hnsw_faiss_n32e40.index")
|
||||
|
||||
# 6. Save chunk list
|
||||
with open("nq_list.pkl", "wb") as f:
|
||||
pickle.dump(chunk_texts, f)
|
||||
|
||||
print("Index and chunk list saved successfully.")
|
||||
```
|
||||
|
||||
|
||||
## Evaluation
|
||||
|
||||
Results are coming soon.
|
||||
|
||||
Reference in New Issue
Block a user