Introduction

Hello there, welcome to the second lesson of our "Scaling Up RAG with Vector Databases" course! In the previous unit, you explored how to break large documents into smaller chunks and attach useful metadata (like doc_id, chunk_id, and labels such as category). These chunks are essential for structuring data in a way that makes retrieval easier. In this lesson, we'll build on that groundwork by showing you how to store them in a vector database. Vector databases are specialized systems designed for high-speed, semantic querying of vectors. By switching from keyword-based searches to semantic searches, your RAG system will retrieve relevant information more efficiently. Let's dive in!

Understanding Vector Databases

A vector database stores data in the form of numerical vectors that capture the semantic essence of texts (or other data). The database then uses similarity metrics — rather than literal word matches — so that conceptually similar items are stored close together. This means searches on vector databases can retrieve contextually relevant results even when keywords are absent. By leveraging approximate or exact nearest-neighbor strategies for similarity, vector databases can scale to handle millions or billions of vectors while still providing quick query responses. This makes them especially suitable for RAG systems, which rely on fast semantic lookups across large collections of text.

Setting Up a Vector Database in Java

Now, let's jump into coding with a vector database in Java. Here's how to set up a vector database client:

import java.io.FileReader;
import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import com.robrua.nlp.bert.Bert;

import tech.amikos.chromadb.Client;
import tech.amikos.chromadb.Collection;
import tech.amikos.chromadb.EmbeddingFunction;
import tech.amikos.chromadb.handler.ApiException;

public class Solution {

    public static void main(String[] args) {
        Client client = new Client(System.getenv("CHROMA_URL"));
        
        // Create a custom embedding function
        EmbeddingFunction embedFunc = new BertSentenceTransformerEmbedding();
        
        // Get or create collection
        Collection collection;
        try {
            // Try to get existing collection
            collection = client.getCollection(collectionName, embedFunc);
        } catch (ApiException e) {
            // Create new collection if it doesn't exist
            collection = client.createCollection(collectionName, null, true, embedFunc);
        }
    }
}

How It Works:

  • Embedding Setup: We use a BERT model to generate vectors for the text chunks. The model maps sentences to a dense vector space, capturing semantic meaning.
  • Client and Collection: We create a Client instance to interact with the ChromaDB and manage collections.
  • Embedding Function: A custom BertSentenceTransformerEmbedding is implemented to handle the conversion between text and vector representations.
Preparing Data and Adding Chunks to the Vector Database

After setting up your client and embedding function, the next step is to prepare your chunks for insertion and add them to the database:

public static List<Map<String, Object>> loadAndChunkDataset(String filePath, int chunkSize) 
        throws IOException {
    Gson gson = new Gson();
    List<Map<String, Object>> data;
    
    try (FileReader reader = new FileReader(filePath)) {
        data = gson.fromJson(reader, new TypeToken<List<Map<String, Object>>>(){}.getType());
    }
    
    List<Map<String, Object>> allChunks = new ArrayList<>();
    
    for (int docId = 0; docId < data.size(); docId++) {
        Map<String, Object> doc = data.get(docId);
        String docText = (String) doc.get("content");
        String docCategory = doc.containsKey("category") ? 
                            (String) doc.get("category") : "general";
        
        List<String> docChunks = chunkText(docText, chunkSize);
        
        for (int chunkId = 0; chunkId < docChunks.size(); chunkId++) {
            Map<String, Object> chunk = new HashMap<>();
            chunk.put("doc_id", docId);
            chunk.put("chunk_id", chunkId);
            chunk.put("category", docCategory);
            chunk.put("text", docChunks.get(chunkId));
            allChunks.add(chunk);
        }
    }
    
    return allChunks;
}

/**
    * Splits the given text into chunks of size 'chunkSize'.
    * Returns a list of chunk strings.
    */
public static List<String> chunkText(String text, int chunkSize) {
    String[] words = text.split("\\s+");
    List<String> chunks = new ArrayList<>();
    
    for (int i = 0; i < words.length; i += chunkSize) {
        int end = Math.min(i + chunkSize, words.length);
        String[] chunkWords = Arrays.copyOfRange(words, i, end);
        chunks.add(String.join(" ", chunkWords));
    }
    
    return chunks;
}

Key Points:

  • Data Grouping: Each chunk is mapped to its text, a unique ID, and metadata. These are used during retrieval and future reference.
  • Seamless Insertion: The prepared data can be inserted into the vector database using Java-compatible methods.
Updating and Managing Documents

Vector databases allow you to keep your collection up to date with new or modified information. Below is an example of adding and then deleting a "document" (or chunk) after the collection has already been created:

public static Collection buildChromaCollection(List<Map<String, Object>> chunks, String collectionName) throws ApiException {
    Client client = new Client(System.getenv("CHROMA_URL"));
    
    // Create a custom embedding function
        EmbeddingFunction embedFunc = new BertSentenceTransformerEmbedding();
        
        // Get or create collection
        Collection collection;
        try {
            // Try to get existing collection
            collection = client.getCollection(collectionName, embedFunc);
        } catch (ApiException e) {
            // Create new collection if it doesn't exist
            collection = client.createCollection(collectionName, null, true, embedFunc);
        }
        
        // Prepare documents, IDs, and metadata
        List<String> texts = new ArrayList<>();
        List<String> ids = new ArrayList<>();
        List<Map<String, String>> metadatas = new ArrayList<>();
        
        for (Map<String, Object> chunk : chunks) {
            int docId = Double.valueOf((String) chunk.get("doc_id")).intValue();
            int chunkId = ((Number) chunk.get("chunk_id")).intValue();
            String category = (String) chunk.get("category");
            String text = (String) chunk.get("text");
            
            String id = "chunk_" + docId + "_" + chunkId;
            
            // Create metadata
            Map<String, String> metadata = new HashMap<>();
            metadata.put("doc_id", Integer.toString(docId));
            metadata.put("chunk_id", Integer.toString(chunkId));
            metadata.put("category", category);
            
            texts.add(text);
            ids.add(id);
            metadatas.add(metadata);
        }
        
        // Add documents to the collection
        collection.add(null,metadatas,texts,ids);
        
        // Delete the added documents
        collection.deleteWithIds(ids);
        
        return collection;
}

Key Points:

  • chunks is our initial list of text chunks with their metadata.
  • id is a unique identifier string created by combining document ID and chunk ID (e.g., "chunk_2_0").
  • Why Unique IDs Matter: Each chunk needs a unique identifier so the vector database can reference it later for updates, deletions, or retrieval. By combining doc_id and chunk_id into a string like "chunk_2_0", we ensure each chunk has a distinct ID while maintaining its relationship to the source document.
  • Adding and Deleting: The prepared data can be added or removed from the vector database using Java-compatible methods. By sending null as the embeddings, they are automatically calculated using the EmbeddingFunction of the collection.
Conclusion and Next Steps

By storing text chunks in a vector database, you've laid the foundation for faster, more semantically aware retrieval. You know how to create, update, and manage a vector database collection — crucial skills for any large-scale RAG system.

In the next lesson, you'll learn how to query the vector database to fetch the most relevant chunks and feed them into a language model. That's where the real magic of producing context-rich, accurate responses shines! For now, feel free to explore different embedding models or try adding and deleting a variety of chunks. When you're ready, proceed to the practice exercises to cement these concepts and further refine your RAG workflow.

Previous
Next
Sign up
Join the 1M+ learners on CodeSignal
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal