Skip to content

Commit 302ab1b

Browse files
authored
[Docs] - Colbert Intro, Example, and Notebook (#401)
* scaffold * Created using Colab * push-notebook * intro * content * struct * more-content * astra-advantages * cleanup * cleanup * adc * vector-index * comparison * rank
1 parent 8bbb1a6 commit 302ab1b

File tree

6 files changed

+3015
-1
lines changed

6 files changed

+3015
-1
lines changed

docs/modules/ROOT/nav.adoc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@
1313
* xref:default-architecture:retrieval.adoc[]
1414
* xref:default-architecture:generation.adoc[]
1515
16+
.ColBERT
17+
* xref:colbert:index.adoc[]
18+
* xref:examples:colbert.adoc[]
19+
1620
.Introduction to RAG
1721
* xref:intro-to-rag:index.adoc[]
1822
* xref:intro-to-rag:indexing.adoc[]
@@ -25,6 +29,7 @@
2529
2630
.RAGStack Examples
2731
* xref:examples:index.adoc[]
32+
* xref:examples:colbert.adoc[]
2833
* xref:examples:langchain_multimodal_gemini.adoc[]
2934
* xref:examples:nvidia_embeddings.adoc[]
3035
* xref:examples:hotels-app.adoc[]

docs/modules/ROOT/pages/packages.adoc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ Additional LLamaIndex packages should work out of the box, although you need to
3636

3737
The `colbert` module provides a vanilla implementation for ColBERT retrieval. It is not tied to any specific framework and can be used with any of the RAGStack packages.
3838

39-
If you want to use ColBERT with LangChain or LLamaIndex, you can use the following the extras:
39+
If you want to use ColBERT with LangChain or LLamaIndex, you can use the following extras:
4040

4141
. `ragstack-ai-langchain[colbert]`
4242
. `ragstack-ai-llamaindex[colbert]`

docs/modules/colbert/pages/index.adoc

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
= Introduction to ColBERT
2+
3+
ColBERT stands for "Contextualized Late Interaction over BERT".
4+
5+
"Contextualized Late Interaction" describes a unique method of interacting with Stanford University's https://arxiv.org/abs/2004.12832[BERT]{external-link-icon} model.
6+
7+
ColBERT is a machine learning retrieval model that improves the computational efficiency and contextual depth of information retrieval tasks.
8+
9+
*TL;DR:*
10+
11+
1. BERT embeds text chunks as matrices of token-level vectors, enabling much deeper context matching than a single vector embedding per chunk.
12+
2. BERT manages this additional depth by pre-processing documents and queries into uniform lengths with the https://huggingface.co/learn/nlp-course/en/chapter6/6[Wordpiece]{external-link-icon} tokenizer, ideal for batch processing on GPUs.
13+
3. "Contextualized Late Interaction" first retrieves the top-k chunks with the highest similarity scores to query tokens.
14+
The top-k chunks are then sorted again. The query tokens are compared to every token in the chunk to rank the chunks by the highest aggregate similarity score.
15+
16+
See the xref:examples:colbert.adoc[ColBERT example code] to get started using ColBERT with RAGStack and Astra DB.
17+
18+
== RAGStack-ai-colbert packages
19+
20+
`ragstack-ai-colbert` contains the implementation of the ColBERT retrieval.
21+
22+
The `colbert` module provides a vanilla implementation for ColBERT retrieval. It is not tied to any specific framework and can be used with any of the RAGStack packages.
23+
24+
To use ColBERT with LangChain or LLamaIndex, install ColBERT as an extra:
25+
26+
* `ragstack-ai-langchain[colbert]`
27+
* `ragstack-ai-llamaindex[colbert]`
28+
29+
== How is ColBERT different from RAG?
30+
31+
In the common RAG usage, a standard embedding model represents each chunk as a single vector embedding.
32+
This is called "sparse embedding".
33+
A cosine similarity search performs similarity matching between document and query embeddings, and the top-k results are returned.
34+
This is fast and straightforward, but some context and efficiency are lost.
35+
36+
In the ColBERT model, each chunk is represented as a list of token-level embedding vectors.
37+
This is called "dense embedding".
38+
This per-token "bag of words" within a chunk offers far deeper context than a single vector per chunk.
39+
Document embeddings are pre-computed and indexed with a uniform length to facilitate batch processing.
40+
41+
ColBERT queries are performed in two stages:
42+
43+
1. The query is embedded (densely) and an Approximate Nearest Neighbor (ANN) search compares every query vector token to every context vector token.
44+
Recall that the BERT context chunks have embeddings for each token, so this is a dense comparison.
45+
The closest matches are returned as the top-k chunks.
46+
2. Contextualized Late Interaction ranks the top-k chunks by a fine-grained similarity score.
47+
For each query’s token embedding, the score function generates a highest similarity score based on the max dot product of the query token vector, and all the token embeddings per chunk. The aggregate of all the max scores across all the query tokens is the overall similarity score of that particular chunk.
48+
49+
A vector index in the database significantly improves the speed of this comparison.
50+
51+
== ColBERT, RAGStack, and Astra DB
52+
53+
The https://huggingface.co/colbert-ir/colbertv2.0[ColBERT v2.0]{external-link-icon} library transforms a text chunk into a matrix of token-level embeddings. The output is a 128-dimensional vector for each token in a chunk. This results in a two-dimensional matrix, which doesn't align with the current LangChain interface that outputs a list of floats.
54+
55+
To solve this problem, the `ragstack-ai-colbert` packages and extras include new classes for mapping token-level embedding vectors to the Astra DB vector database.
56+
57+
The https://github.com/datastax/ragstack-ai/blob/main/libs/colbert/ragstack_colbert/cassandra_vector_store.py#L20C7-L20C27[CassandraVectorStore]{external-link-icon} class extends the `BaseVectorStore` class to store and retrieve vector embeddings generated by ColBERT.
58+
59+
The Contextualized Late Interaction retrieval logic is defined in the https://github.com/datastax/ragstack-ai/blob/main/libs/colbert/ragstack_colbert/colbert_retriever.py[ColbertRetriever]{external-link-icon} class, which asynchronously retrieves and scores chunks.
60+
61+
Together, these classes enable ColBERT to be used with Astra DB and the RAGStack ecosystem. But why expend all this effort to use ColBERT with Astra DB?
62+
63+
== Advantages of ColBERT on Astra DB
64+
65+
Our testing with ColBERT has shown that it delivers significantly better recall than any single-vector encodings, but this comes at the cost of a significantly larger dataset size.
66+
67+
For a dataset with 10 million passages (which is ~25% of the English-language Wikipedia), the OpenAI-v3-small model requires 61.44GB, while the ColBERT model requires 768GB.
68+
69+
Most vector indexes can't scale to this size due to problems with *index segmentation* and *memory footprint*.
70+
Astra DB incorporates new vector indexing techniques to address these issues.
71+
72+
*Index segmentation* problems degrade query time. Most vector databases can't index data larger than available RAM in a single physical index, so larger logical indexes are created by splitting the dataset into memory-sized segments. The problem with this approach is that searching within a segment is a logarithmic-time operation in the number of vectors, while combining results across multiple segments is linear-time. So, as your data set grows past the maximum size of a single segment, query time quickly degrades.
73+
74+
Astra DB has larger-than-memory index construction that allows over an order of magnitude more vectors in a single index segment.
75+
76+
*Memory footprint* problems are expensive. Most vector databases require memory proportional to the number of vectors to serve requests. This is done with either the original full-resolution vectors, or quantized (compressed) vectors. But even 16x compression (which must be done with appropriate reranking to avoid destroying recall) requires 48GB of RAM dedicated to just the compressed vectors of a ColBERT dataset of 10M passages (about 1.5B vectors). Adding in other indexes, (graph index) edge caching, and row caching can easily require expensive 128GB server instances.
77+
78+
Astra DB has fused Asymmetric Distance Computation (ADC) graph traversal that reduces the in-memory footprint of a vector index to near-zero.
79+
80+
On top of these improvements, Astra DB preserves the nonblocking index structure and synchronous, realtime index updates that are the hallmark of Astra’s non-Vector indexes.
81+
82+
DataStax has open-sourced the underlying index technology as https://github.com/jbellis/jvector/[JVector]{external-link-icon}.
83+
84+
For more on the challenges of vector indexing at scale, see:
85+
86+
* https://stackoverflow.com/questions/2703432/what-are-segments-in-lucene[Segments in Lucene]{external-link-icon}
87+
* https://thenewstack.io/why-vector-size-matters/[Why Vector Size Matters]{external-link-icon}
Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
= ColBERT in RAGStack with Astra
2+
3+
image::https://colab.research.google.com/assets/colab-badge.svg[align="left",link="https://colab.research.google.com/github/datastax/ragstack-ai/blob/main/examples/notebooks/RAGStackColBERT.ipynb"]
4+
5+
Use ColBERT, Astra DB, and RAGStack to:
6+
7+
. Create ColBERT embeddings
8+
. Index embeddings on Astra
9+
. Retrieve embeddings with RAGStack and Astra
10+
. Use the LangChain ColBERT retriever plugin
11+
12+
== Prerequisites
13+
14+
Import the ragstack-ai-colbert package:
15+
[source,python]
16+
----
17+
pip install ragstack-ai-colbert
18+
----
19+
20+
== Prepare data and create embeddings
21+
22+
. Prepare documents for chunking.
23+
+
24+
[source,python]
25+
----
26+
arctic_botany_dict = {
27+
"Introduction to Arctic Botany": "Arctic botany is the study of plant life in the Arctic, a region characterized by extreme cold, permafrost, and minimal sunlight for much of the year. Despite these harsh conditions, a diverse range of flora thrives here, adapted to survive with minimal water, low temperatures, and high light levels during the summer. This introduction aims to shed light on the resilience and adaptation of Arctic plants, setting the stage for a deeper dive into the unique botanical ecosystem of the Arctic.",
28+
"Arctic Plant Adaptations": "Plants in the Arctic have developed unique adaptations to endure the extreme climate. Perennial growth, antifreeze proteins, and a short growth cycle are among the evolutionary solutions. These adaptations not only allow the plants to survive but also to reproduce in short summer months. Arctic plants often have small, dark leaves to absorb maximum sunlight, and some species grow in cushion or mat forms to resist cold winds. Understanding these adaptations provides insights into the resilience of Arctic flora.",
29+
"The Tundra Biome": "The Arctic tundra is a vast, treeless biome where the subsoil is permanently frozen. Here, the vegetation is predominantly composed of dwarf shrubs, grasses, mosses, and lichens. The tundra supports a surprisingly rich biodiversity, adapted to its cold, dry, and windy conditions. The biome plays a crucial role in the Earth's climate system, acting as a carbon sink. However, it's sensitive to climate change, with thawing permafrost and shifting vegetation patterns.",
30+
"Arctic Plant Biodiversity": "Despite the challenging environment, the Arctic boasts a significant variety of plant species, each adapted to its niche. From the colorful blooms of Arctic poppies to the hardy dwarf willows, these plants form a complex ecosystem. The biodiversity of Arctic flora is vital for local wildlife, providing food and habitat. This diversity also has implications for Arctic peoples, who depend on certain plant species for food, medicine, and materials.",
31+
"Climate Change and Arctic Flora": "Climate change poses a significant threat to Arctic botany, with rising temperatures, melting permafrost, and changing precipitation patterns. These changes can lead to shifts in plant distribution, phenology, and the composition of the Arctic flora. Some species may thrive, while others could face extinction. This dynamic is critical to understanding future Arctic ecosystems and their global impact, including feedback loops that may exacerbate global warming.",
32+
"Research and Conservation in the Arctic": "Research in Arctic botany is crucial for understanding the intricate balance of this ecosystem and the impacts of climate change. Scientists conduct studies on plant physiology, genetics, and ecosystem dynamics. Conservation efforts are focused on protecting the Arctic's unique biodiversity through protected areas, sustainable management practices, and international cooperation. These efforts aim to preserve the Arctic flora for future generations and maintain its role in the global climate system.",
33+
"Traditional Knowledge and Arctic Botany": "Indigenous peoples of the Arctic have a deep connection with the land and its plant life. Traditional knowledge, passed down through generations, includes the uses of plants for nutrition, healing, and materials. This body of knowledge is invaluable for both conservation and understanding the ecological relationships in Arctic ecosystems. Integrating traditional knowledge with scientific research enriches our comprehension of Arctic botany and enhances conservation strategies.",
34+
"Future Directions in Arctic Botanical Studies": "The future of Arctic botany lies in interdisciplinary research, combining traditional knowledge with modern scientific techniques. As the Arctic undergoes rapid changes, understanding the ecological, cultural, and climatic dimensions of Arctic flora becomes increasingly important. Future research will need to address the challenges of climate change, explore the potential for Arctic plants in biotechnology, and continue to conserve this unique biome. The resilience of Arctic flora offers lessons in adaptation and survival relevant to global challenges."
35+
}
36+
arctic_botany_chunks = list(arctic_botany_dict.values())
37+
----
38+
+
39+
. Start the ColBERT configuration and create embeddings.
40+
+
41+
[source,python]
42+
----
43+
from ragstack_colbert import ColbertEmbeddingModel, ChunkData
44+
# colbert stuff starts
45+
colbert = ColbertEmbeddingModel()
46+
47+
chunks = [ChunkData(text=text, metadata={}) for text in arctic_botany_chunks]
48+
49+
embedded_chunks = colbert.embed_chunks(chunks=chunks, doc_id="arctic botany")
50+
----
51+
+
52+
. Examine the embeddings.
53+
+
54+
[source,python]
55+
----
56+
assert len(embedded_chunks) == 8
57+
----
58+
59+
== Create a vector store in Astra
60+
61+
. Ingest embeddings and create a vector store in Astra.
62+
+
63+
[source,python]
64+
----
65+
from ragstack_colbert import CassandraVectorStore
66+
from getpass import getpass
67+
import cassio
68+
69+
keyspace="default_keyspace"
70+
database_id=getpass("Enter your Astra Database Id:")
71+
astra_token=getpass("Enter your Astra Token:")
72+
73+
cassio.init(token=astra_token, database_id=database_id, keyspace=keyspace)
74+
session=cassio.config.resolve_session()
75+
76+
db = CassandraVectorStore(
77+
keyspace = keyspace,
78+
table_name="colbert_embeddings4",
79+
session = session,
80+
)
81+
----
82+
+
83+
. Create an index in Astra.
84+
+
85+
[source,python]
86+
----
87+
db.put_chunks(chunks=embedded_chunks, delete_existing=True)
88+
----
89+
90+
== Retrieve embeddings from the Astra index
91+
92+
Create a RAGStack retriever and ask questions against the indexed embeddings.
93+
The library provides:
94+
95+
* Embed query tokens
96+
* Generate candidate documents using Astra ANN search
97+
* Max similarity scoring
98+
* Ranking
99+
100+
[source,python]
101+
----
102+
import logging
103+
import nest_asyncio
104+
nest_asyncio.apply()
105+
106+
logging.getLogger('cassandra').setLevel(logging.ERROR) # workaround to suppress logs
107+
from ragstack_colbert import ColbertRetriever
108+
retriever = ColbertRetriever(
109+
vector_store=db, embedding_model=colbert
110+
)
111+
112+
answers = retriever.retrieve("What's artic botany", k=2)
113+
for answer in answers:
114+
print(f"Rank: {answer.rank} Score: {answer.score} Text: {answer.data.text}\n")
115+
----
116+
117+
.Result
118+
[source, plain]
119+
----
120+
Rank: 1 Score: 35.225955963134766 Text: Arctic botany is the study of plant life in the Arctic, a region characterized by extreme cold, permafrost, and minimal sunlight for much of the year. Despite these harsh conditions, a diverse range of flora thrives here, adapted to survive with minimal water, low temperatures, and high light levels during the summer. This introduction aims to shed light on the resilience and adaptation of Arctic plants, setting the stage for a deeper dive into the unique botanical ecosystem of the Arctic.
121+
122+
Rank: 2 Score: 29.655662536621094 Text: Research in Arctic botany is crucial for understanding the intricate balance of this ecosystem and the impacts of climate change. Scientists conduct studies on plant physiology, genetics, and ecosystem dynamics. Conservation efforts are focused on protecting the Arctic's unique biodiversity through protected areas, sustainable management practices, and international cooperation. These efforts aim to preserve the Arctic flora for future generations and maintain its role in the global climate system.
123+
----
124+
125+
== LangChain retriever
126+
127+
Alternatively, use the ColBERT extra with your RAGStack package to retrieve documents.
128+
129+
. Install the RAGStack LangChain package with the ColBERT extra.
130+
+
131+
[source,python]
132+
----
133+
pip install ragstack-ai-langchain[colbert]
134+
----
135+
+
136+
. Run the LangChain retriever against the indexed embeddings.
137+
+
138+
[source,python]
139+
----
140+
from ragstack_langchain.colbert import ColbertLCRetriever
141+
142+
lc_retriever = ColbertLCRetriever(retriever, k=2)
143+
docs = lc_retriever.get_relevant_documents("what kind fish lives shallow coral reefs atlantic, india ocean, red sea, gulf of mexico, pacific, and arctic ocean")
144+
print(f"first answer: {docs[0].page_content}")
145+
----
146+
+
147+
.Result
148+
[source,plain]
149+
----
150+
....
151+
first answer: Despite the challenging environment, the Arctic boasts a significant variety of plant species, each adapted to its niche. From the colorful blooms of Arctic poppies to the hardy dwarf willows, these plants form a complex ecosystem. The biodiversity of Arctic flora is vital for local wildlife, providing food and habitat. This diversity also has implications for Arctic peoples, who depend on certain plant species for food, medicine, and materials.
152+
....
153+
----

docs/modules/examples/pages/index.adoc

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,10 @@ a| image::https://colab.research.google.com/assets/colab-badge.svg[align="left",
7171
|===
7272
| Description | Colab | Documentation
7373

74+
| Create ColBERT embeddings, index embeddings on Astra, and retrieve embeddings with RAGStack.
75+
a| image::https://colab.research.google.com/assets/colab-badge.svg[align="left",link="https://colab.research.google.com/github/datastax/ragstack-ai/blob/main/examples/notebooks/RAGStackColBERT.ipynb"]
76+
| xref:colbert.adoc[]
77+
7478
| Implement a generative Q&A over your own documentation with {db-serverless} Search, OpenAI, and CassIO.
7579
a| image::https://colab.research.google.com/assets/colab-badge.svg[align="left",link="https://colab.research.google.com/github/datastax/ragstack-ai/blob/main/examples/notebooks/QA_with_cassio.ipynb"]
7680
| xref:qa-with-cassio.adoc[]

0 commit comments

Comments
 (0)