Storing and Managing Text Chunks in a Vector Database with Java

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);
        }
    }
}

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;
}

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;
}