ColBERT in RAGStack with Astra

Use ColBERT, Astra DB, and RAGStack to:

Create ColBERT embeddings.
Index embeddings on Astra.
Retrieve embeddings with RAGStack.
1. Optionally, use the LangChain ColBERT retriever plugin.

For more information, see Introduction to ColBERT.

To run ragstack-ai-colbert in a Windows environment, use Windows Subsystem for Linux.

Prerequisites

Install dependencies:

pip install ragstack-ai-colbert python-dotenv

Create a .env file in your application directory with the following environment variables:
```
ASTRA_DB_APPLICATION_TOKEN=AstraCS: ...
ASTRA_DB_ID=2eab82dc-9032-45ba-aeb0-a481b6f9458d
```
In an Astra API endpoint like https://2eab82dc-9032-45ba-aeb0-a481b6f9458d-us-east-1.apps.astra.datastax.com, the ASTRA_DB_ID is 2eab82dc-9032-45ba-aeb0-a481b6f9458d.

Prepare data and create embeddings

Import dependencies and load environment variables.

import os
import logging
import nest_asyncio
from dotenv import load_dotenv
from ragstack_colbert import CassandraDatabase, ColbertEmbeddingModel, ColbertVectorStore

load_dotenv()

Set up the ColBERT and Astra configurations.

keyspace="default_keyspace"
database_id=os.getenv("ASTRA_DB_ID")
astra_token=os.getenv("ASTRA_DB_APPLICATION_TOKEN")

database = CassandraDatabase.from_astra(
    astra_token=astra_token,
    database_id=database_id,
    keyspace=keyspace
)

embedding_model = ColbertEmbeddingModel()

vector_store = ColbertVectorStore(
    database = database,
    embedding_model = embedding_model,
)

Prepare documents for chunking.

arctic_botany_dict = {
    "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.",
    "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.",
    "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.",
    "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.",
    "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.",
    "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.",
    "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.",
    "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."
}

arctic_botany_texts = list(arctic_botany_dict.values())

Connect to Astra and ingest embeddings.

results = vector_store.add_texts(texts=arctic_botany_texts, doc_id="arctic_botany")

Retrieve embeddings from the Astra index

Create a RAGStack retriever and query the indexed embeddings. The library includes:

Query token embedding
Candidate documents generated using Astra ANN search
Max similarity scoring

Ranking

Python
Result

nest_asyncio.apply()

logging.getLogger("cassandra").setLevel(logging.ERROR)  # workaround to suppress logs
retriever = vector_store.as_retriever()

answers = retriever.text_search("What's arctic botany", k=2)
for rank, (answer, score) in enumerate(answers):
    print(f"Rank: {rank} Score: {score} Text: {answer.text}\n")

#> QueryTokenizer.tensorize(batch_text[0], batch_background[0], bsize) ==
#> Input: . What's arctic botany,                 True,           None
#> Output IDs: torch.Size([9]), tensor([  101,     1,  2054,  1005,  1055,  2396,  2594, 17018,   102])
#> Output Mask: torch.Size([9]), tensor([1, 1, 1, 1, 1, 1, 1, 1, 1])

Rank: 0 Score: 5.266004428267479 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.

Rank: 1 Score: 5.266004309058189 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.

Retrieve embeddings with the LangChain retriever

Alternatively, use the ColBERT extra with the ragstack-ai-langchain package to retrieve documents.

Install the RAGStack Langchain package with the ColBERT extra.
```
pip install "ragstack-ai-langchain[colbert]"
```

Run the LangChain retriever against the indexed embeddings.

Python
Result

from ragstack_langchain.colbert import ColbertVectorStore as LangchainColbertVectorStore

lc_vector_store = LangchainColbertVectorStore(
    database=database,
    embedding_model=embedding_model,
)

docs = lc_vector_store.similarity_search(
    "what kind fish lives shallow coral reefs atlantic, india ocean, "
    "red sea, gulf of mexico, pacific, and arctic ocean"
)
print(f"first answer: {docs[0].page_content}")

....
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.
....

ColBERT in RAGStack with Astra

Prerequisites

Prepare data and create embeddings

Retrieve embeddings from the Astra index

Retrieve embeddings with the LangChain retriever

Was this helpful?

Give Feedback