from dashtext import TextTokenizer, SparseVectorEncoder

my_tokenizer = TextTokenizer.from_pretrained(model_name='Jieba', dict='dict.txt')
my_encoder = SparseVectorEncoder(tokenize_function=my_tokenizer.tokenize)

from dashtext import SparseVectorEncoder
from transformers import BertTokenizer

my_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
my_encoder = SparseVectorEncoder(tokenize_function=my_tokenizer.tokenize)

import com.aliyun.dashtext.common.DashTextException;
import com.aliyun.dashtext.common.ErrorCode;
import com.aliyun.dashtext.encoder.SparseVectorEncoder;
import com.aliyun.dashtext.tokenizer.BaseTokenizer;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

public class Main {
    public static class MyTokenizer implements BaseTokenizer {
        @Override
        public List<String> tokenize(String s) throws DashTextException {
            if (s == null) {
                throw new DashTextException(ErrorCode.INVALID_ARGUMENT);
            }

            // 使用正則表達(dá)式將文本按空白符和標(biāo)點(diǎn)符號(hào)分割，并轉(zhuǎn)換為小寫
            return Arrays.stream(s.split("\\s+|(?<!\\d)[.,](?!\\d)"))
                    .map(String::toLowerCase)
                    .filter(token -> !token.isEmpty())   // 過(guò)濾掉空字符串
                    .collect(Collectors.toList());
        }
    }

    public static void main(String[] args) {
        SparseVectorEncoder encoder = new SparseVectorEncoder(new MyTokenizer());
    }
}

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


# 創(chuàng)建空SparseVectorEncoder(可以設(shè)置自定義Tokenizer)
encoder = SparseVectorEncoder()


# step1: 準(zhǔn)備語(yǔ)料以及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",
]


# step2: 訓(xùn)練BM25 Model
encoder.train(corpus_document)


# step3: 調(diào)參優(yōu)化BM25 Model
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: ")
    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)


# step4: 選擇最優(yōu)參數(shù)并保存模型
encoder.b = 0.75
encoder.k1 = 1.5
encoder.dump("./model.json")


# step5: 后續(xù)使用時(shí)可以加載模型
new_encoder = SparseVectorEncoder()
bm25_model_path = "./model.json"
new_encoder.load(bm25_model_path)


# step6: 對(duì)模型進(jìn)行finetune并保存
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 {
        // 創(chuàng)建空SparseVectorEncoder(可以設(shè)置自定義Tokenizer)
        SparseVectorEncoder encoder = new SparseVectorEncoder();

        // step1: 準(zhǔn)備語(yǔ)料以及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"
        );

        // step2: 訓(xùn)練BM25 Model
        encoder.train(corpusDocument);

        // step3: 調(diào)參優(yōu)化BM25 Model
        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);
            }
        }

        // step4: 選擇最優(yōu)參數(shù)并保存模型
        encoder.setB(0.75f);
        encoder.setK1(1.5f);
        encoder.dump("./model.json");

        // step5: 后續(xù)使用時(shí)可以加載模型
        SparseVectorEncoder newEncoder = new SparseVectorEncoder();
        newEncoder.load("./model.json");

        // step6: 對(duì)模型進(jìn)行finetune并保存
        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");
    }
}