Skip to content

Latest commit

 

History

History
133 lines (100 loc) · 3.19 KB

README.md

File metadata and controls

133 lines (100 loc) · 3.19 KB

Colbert

This is a playground to use ColBERT It does not use RAGatouille

The problem with RAGatouille are

  • It does not expose all the ColBERT configurations.
  • It is difficult directly integrate with LangChain's Embeddings class

Gotcha

The latest version of colbert-ai==0.2.19 or its dependencies require pyarraow==14.0.0

Install faiss-gpu on CUDA

what's in the repo

Code is at this folder that includes

  • A ColBERT Embedding class
  • Astra loader
  • Astra vector based retriever, a LangChain compatible retriever
  • It runs on CPU and GPU/Cuda (automatically runs all available GPUs) A chat bot example of RAG using ColBERT embedding, Astra DB vector store, retriever (including a default ranker).

How to run the example and prerequisites:

  • Specify the directory of pdf files
  • Create a AstraDB keyspace and specify the keyspace name in the example code
  • Download Secure Connect Bundle and specify the path in the example
  • Create an AstraCS token to export as ASTRA_TOKEN
cd webserver
poetry install
poetry shell
cd webserver
python example.py

extra

  • A web server for embedding service
  • Dockerfile of the web embedding service
  • Indexing and encoding examples[example] to test on GPU.

Examples

LangChain/RagSTACK loader and splitter

Load, split and prepare the documents

from langchain_community.document_loaders import DirectoryLoader
from langchain_community.document_loaders import PyPDFLoader
import os

# pip install pypdf
loader =DirectoryLoader(
    path="./files",
    glob="**/*.pdf",
    loader_cls=PyPDFLoader,
    recursive=True,
)

docs = loader.load()

from langchain.text_splitter import CharacterTextSplitter
text_splitter = CharacterTextSplitter(
    separator="\n",
    chunk_size=500, # colbert doc_maxlen is 220
    chunk_overlap=100,
    length_function=len,
)

splits = text_splitter.split_documents(docs)
title = docs[0].metadata['source']
collections = []

for part in splits:
    collections.append(part.page_content)

Create ColBERT Embeddings

from embedding import ColbertTokenEmbeddings

colbert = ColbertTokenEmbeddings(
    doc_maxlen=220,
    nbits=1,
    kmeans_niters=4,
    nranks=1,
)

passageEmbeddings = colbert.embed_documents(texts=collections, title=title)

Load ColBERT embeddings into Astra DB

Create tables and load embeddings

from embedding import AstraDB
import os

# astra db
astra = AstraDB(
    secure_connect_bundle="./secure-connect-mingv1.zip",
    astra_token=os.getenv("ASTRA_TOKEN"),
    keyspace="colbert128"
)

Retrieval from AstraDB

from embedding import ColbertAstraRetriever

retriever = ColbertAstraRetriever(astraDB=astra, colbertEmbeddings=colbert)
answers = retriever.retrieve("what's the toll free number to call for help?")

Web embedding service

A web embedding service is implemented to provide ColBERT text embedding over HTTP.

Commands to set up dev environment.

cd webserver
poetry install
poetry shell
cd webserver
uvicorn main:app --reload

Performance and configuration