Advanced use

from dashtext import SparseVectorEncoder
from pydantic import BaseModel
from typing import Dict, List


class Result(BaseModel):
    doc: str
    score: float


def calculate_score(query_vector: Dict[int, float], document_vector: Dict[int, float]) -> float:
    score = 0.0
    for key, value in query_vector.items():
        if key in document_vector:
            score += value * document_vector[key]
    return score


# Create an empty SparseVectorEncoder. You can set a custom tokenizer.
encoder = SparseVectorEncoder()


# Step 1. Prepare the corpus and documents
corpus_document: List[str] = [
    "The quick brown fox rapidly and agilely leaps over the lazy dog that lies idly by the roadside.",
    "Never jump over the lazy dog quickly",
    "A fox is quick and jumps over dogs",
    "The quick brown fox",
    "Dogs are domestic animals",
    "Some dog breeds are quick and jump high",
    "Foxes are wild animals and often have a brown coat",
]


# Step 2. Train the BM25 model
encoder.train(corpus_document)


# Step 3. Fine-tune the parameters
query: str = "quick brown fox"
print(f"query: {query}")
k1s = [1.0, 1.5]
bs = [0.5, 0.75]
for k1, b in zip(k1s, bs):
    print(f"current k1: {k1}, b: {b}")
    encoder.b = b
    encoder.k1 = k1
    query_vector = encoder.encode_queries(query)
    results: List[Result] = []
    for idx, doc in enumerate(corpus_document):
        doc_vector = encoder.encode_documents(doc)
        score = calculate_score(query_vector, doc_vector)
        results.append(Result(doc=doc, score=score))
    results.sort(key=lambda r: r.score, reverse=True)

    for result in results:
        print(result)


# Step 4. Select the optimal parameters and save the model
encoder.b = 0.75
encoder.k1 = 1.5
encoder.dump("./model.json")


# Step 5. Load the model for subsequent use
new_encoder = SparseVectorEncoder()
bm25_model_path = "./model.json"
new_encoder.load(bm25_model_path)


# Step 6. Incremental update: add new corpus data and save the updated model
extra_corpus: List[str] = [
    "The fast fox jumps over the lazy, chubby dog",
    "A swift fox hops over a napping old dog",
    "The quick fox leaps over the sleepy, plump dog",
    "The agile fox jumps over the dozing, heavy-set dog",
    "A speedy fox vaults over a lazy, old dog lying in the sun"
]

new_encoder.train(extra_corpus)
new_bm25_model_path = "new_model.json"
new_encoder.dump(new_bm25_model_path)

import com.aliyun.dashtext.encoder.SparseVectorEncoder;

import java.io.*;
import java.util.*;

public class Main {

    public static class Result {
        public String doc;
        public float score;

        public Result(String doc, float score) {
            this.doc = doc;
            this.score = score;
        }

        @Override
        public String toString() {
            return String.format("Result(doc=%s, score=%f)", doc, score);
        }
    }

    public static float calculateScore(Map<Long, Float> queryVector, Map<Long, Float> documentVector) {
        float score = 0.0f;
        for (Map.Entry<Long, Float> entry : queryVector.entrySet()) {
            if (documentVector.containsKey(entry.getKey())) {
                score += entry.getValue() * documentVector.get(entry.getKey());
            }
        }
        return score;
    }

    public static void main(String[] args) throws IOException {
        // Create an empty SparseVectorEncoder. You can set a custom tokenizer.
        SparseVectorEncoder encoder = new SparseVectorEncoder();

        // Step 1. Prepare the corpus and documents
        List<String> corpusDocument = Arrays.asList(
                "The quick brown fox rapidly and agilely leaps over the lazy dog that lies idly by the roadside.",
                "Never jump over the lazy dog quickly",
                "A fox is quick and jumps over dogs",
                "The quick brown fox",
                "Dogs are domestic animals",
                "Some dog breeds are quick and jump high",
                "Foxes are wild animals and often have a brown coat"
        );

        // Step 2. Train the BM25 model
        encoder.train(corpusDocument);

        // Step 3. Fine-tune the parameters
        String query = "quick brown fox";
        System.out.println("query: " + query);
        float[] k1s = {1.0f, 1.5f};
        float[] bs = {0.5f, 0.75f};
        for (int i = 0; i < k1s.length; i++) {
            float k1 = k1s[i];
            float b = bs[i];
            System.out.println("current k1: " + k1 + ", b: " + b);
            encoder.setB(b);
            encoder.setK1(k1);

            Map<Long, Float> queryVector = encoder.encodeQueries(query);
            List<Result> results = new ArrayList<>();
            for (String doc : corpusDocument) {
                Map<Long, Float> docVector = encoder.encodeDocuments(doc);
                float score = calculateScore(queryVector, docVector);
                results.add(new Result(doc, score));
            }

            results.sort((r1, r2) -> Float.compare(r2.score, r1.score));

            for (Result result : results) {
                System.out.println(result);
            }
        }

        // Step 4. Select the optimal parameters and save the model
        encoder.setB(0.75f);
        encoder.setK1(1.5f);
        encoder.dump("./model.json");

        // Step 5. Load the model for subsequent use
        SparseVectorEncoder newEncoder = new SparseVectorEncoder();
        newEncoder.load("./model.json");

        // Step 6. Incremental update: add new corpus data and save the updated model
        List<String> extraCorpus = Arrays.asList(
                "The fast fox jumps over the lazy, chubby dog",
                "A swift fox hops over a napping old dog",
                "The quick fox leaps over the sleepy, plump dog",
                "The agile fox jumps over the dozing, heavy-set dog",
                "A speedy fox vaults over a lazy, old dog lying in the sun"
        );
        newEncoder.train(extraCorpus);
        newEncoder.dump("./new_model.json");
    }
}