Skip to content

Commit d52c1cf

Browse files
kdroidFilterkdroidFilter
committed
feat(search): dense semantic search with v5 embedding model
- Add SeforimEmbedder (ONNX v5 int8) + HebrewV5Normalizer (final-letter folding) - HybridSearchEngine (BM25 + dense, RRF) + VectorSearcher over a fused Lucene index - BuildVectorIndex / fused KnnFloatVectorField indexing in the generator - Bundle + fetch the v5 model (PackageArtifacts, DownloadEmbedModel -> v5-int8) - CI: free disk space on the runner before the build
1 parent 3ac2df4 commit d52c1cf

15 files changed

Lines changed: 864 additions & 5 deletions

File tree

‎generator/packaging/build.gradle.kts‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,12 +80,37 @@ tasks.register<JavaExec>("downloadLexicalDb") {
8080
jvmArgs = listOf("-Xmx512m")
8181
}
8282

83+
// Download the dense embedding model (int8 ONNX + tokenizer) from the private
84+
// SeforimEmbedding v4-int8 release next to seforim.db, so it gets bundled.
85+
// Needs GITHUB_TOKEN/GH_TOKEN; fails soft (bundle without model) if unavailable.
86+
// Usage:
87+
// GITHUB_TOKEN=… ./gradlew :packaging:downloadEmbedModel
88+
tasks.register<JavaExec>("downloadEmbedModel") {
89+
group = "application"
90+
description = "Download int8 embedding model + tokenizer from the private v4-int8 release next to seforim.db."
91+
92+
dependsOn("jvmJar")
93+
mainClass.set("io.github.kdroidfilter.seforimlibrary.packaging.DownloadEmbedModelKt")
94+
classpath = files(tasks.named("jvmJar")) + configurations.getByName("jvmRuntimeClasspath")
95+
96+
if (project.hasProperty("seforimDb")) {
97+
systemProperty("seforimDb", project.property("seforimDb") as String)
98+
} else if (System.getenv("SEFORIM_DB") != null) {
99+
systemProperty("seforimDb", System.getenv("SEFORIM_DB"))
100+
} else {
101+
val defaultDbPath = rootProject.layout.buildDirectory.file("seforim.db").get().asFile.absolutePath
102+
systemProperty("seforimDb", defaultDbPath)
103+
}
104+
105+
jvmArgs = listOf("-Xmx256m")
106+
}
107+
83108
// Package DB + Lucene indexes into single tar.zst and split
84109
tasks.register<JavaExec>("packageArtifacts") {
85110
group = "application"
86111
description = "Create seforim_bundle.tar.zst (DB + indexes + release info) with zstd and split into ~1.9GiB parts."
87112

88-
dependsOn("jvmJar", "writeReleaseInfo", "downloadLexicalDb")
113+
dependsOn("jvmJar", "writeReleaseInfo", "downloadLexicalDb", "downloadEmbedModel")
89114
mainClass.set("io.github.kdroidfilter.seforimlibrary.packaging.PackageArtifactsKt")
90115
classpath = files(tasks.named("jvmJar")) + configurations.getByName("jvmRuntimeClasspath")
91116

‎generator/packaging/src/jvmMain/kotlin/io/github/kdroidfilter/seforimlibrary/packaging/DownloadEmbedModel.kt‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package io.github.kdroidfilter.seforimlibrary.packaging
2+
3+
import co.touchlab.kermit.Logger
4+
import co.touchlab.kermit.Severity
5+
import io.github.kdroidfilter.seforimlibrary.common.OptimizedHttpClient
6+
import java.nio.file.Files
7+
import java.nio.file.Path
8+
import java.nio.file.Paths
9+
10+
private const val RELEASE_API =
11+
"https://api.github.com/repos/kdroidFilter/SeforimEmbedding/releases/tags/v5-int8"
12+
private const val USER_AGENT = "SeforimLibrary-DownloadEmbedModel/1.0"
13+
14+
// Runtime dense-search artifacts pulled from the private SeforimEmbedding release.
15+
private val ASSETS = listOf("seforim-embed-v5-int8.onnx", "tokenizer.json")
16+
17+
/**
18+
* Download the int8 embedding model + tokenizer from the private `v5-int8` release
19+
* and place them next to `seforim.db` so [PackageArtifacts] bundles them.
20+
*
21+
* Requires a token with read access to the private repo via `GITHUB_TOKEN` / `GH_TOKEN`
22+
* (consumed by [OptimizedHttpClient]). On any failure (no token, network, missing
23+
* asset) it logs a warning and exits 0 so packaging proceeds WITHOUT the model
24+
* (the app then degrades to lexical-only search).
25+
*
26+
* Usage:
27+
* ./gradlew :packaging:downloadEmbedModel
28+
* ./gradlew :packaging:downloadEmbedModel -PseforimDb=/path/to/seforim.db
29+
*/
30+
fun main(args: Array<String>) {
31+
Logger.setMinSeverity(Severity.Info)
32+
val logger = Logger.withTag("DownloadEmbedModel")
33+
34+
val dbPath = resolveDbPath(args)
35+
36+
val present = ASSETS.all { name ->
37+
val p = dbPath.resolveSibling(name)
38+
Files.exists(p) && Files.isRegularFile(p) && Files.size(p) > 0
39+
}
40+
if (present) {
41+
logger.i { "Embedding model already present next to ${dbPath.fileName}; skipping download" }
42+
return
43+
}
44+
45+
runCatching {
46+
val json = OptimizedHttpClient.fetchJson(RELEASE_API, USER_AGENT, logger)
47+
for (name in ASSETS) {
48+
val out = dbPath.resolveSibling(name)
49+
if (Files.exists(out) && Files.size(out) > 0) {
50+
logger.i { "Using existing $name" }
51+
continue
52+
}
53+
val url = assetApiUrl(json, name)
54+
?: throw IllegalStateException("Asset '$name' not found in v4-int8 release")
55+
Files.createDirectories(out.parent)
56+
val tmp = out.resolveSibling("${out.fileName}.part")
57+
Files.deleteIfExists(tmp)
58+
// Asset API url + Accept: octet-stream + token -> works for private repos.
59+
OptimizedHttpClient.downloadFile(url, tmp, USER_AGENT, logger, "Downloading $name")
60+
Files.deleteIfExists(out)
61+
Files.move(tmp, out)
62+
logger.i { "Downloaded $name -> ${out.toAbsolutePath()}" }
63+
}
64+
}.onFailure {
65+
logger.w(it) { "Could not download the embedding model; bundle will omit it (dense search disabled)" }
66+
}
67+
}
68+
69+
private fun resolveDbPath(args: Array<String>): Path {
70+
val dbPathStr = args.getOrNull(0)
71+
?: System.getProperty("seforimDb")
72+
?: System.getenv("SEFORIM_DB")
73+
?: Paths.get("build", "seforim.db").toString()
74+
return Paths.get(dbPathStr)
75+
}
76+
77+
/** Extract the GitHub *asset API* url for a given asset name from the release JSON. */
78+
private fun assetApiUrl(json: String, name: String): String? {
79+
val re = Regex(
80+
"\\{\"url\":\"(https://api\\.github\\.com/[^\"]+?/assets/\\d+)\"[^{}]*?\"name\":\"" +
81+
Regex.escape(name) + "\"",
82+
)
83+
return re.find(json)?.groupValues?.get(1)
84+
}

‎generator/packaging/src/jvmMain/kotlin/io/github/kdroidfilter/seforimlibrary/packaging/PackageArtifacts.kt‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,13 +65,18 @@ fun main(args: Array<String>) {
6565

6666
// Resolve precomputed catalog next to the DB
6767
val catalogPath: Path = dbPath.resolveSibling("catalog.pb")
68-
68+
6969
// Resolve release info file next to the DB
7070
val releaseInfoPath: Path = dbPath.resolveSibling("release_info.txt")
7171

7272
// Resolve lexical DB next to the DB
7373
val lexicalDbPath: Path = dbPath.resolveSibling("lexical.db")
7474

75+
// Resolve the dense embedding model (int8 ONNX) + tokenizer next to the DB.
76+
// Bundled so the app gets dense search out of the box; absent -> lexical only.
77+
val embedModelPath: Path = dbPath.resolveSibling("seforim-embed-v5-int8.onnx")
78+
val embedTokenizerPath: Path = dbPath.resolveSibling("tokenizer.json")
79+
7580
if (!textIndexDir.toFile().isDirectory) {
7681
logger.w { "Lucene text index directory missing: $textIndexDir (will skip)" }
7782
}
@@ -126,6 +131,8 @@ fun main(args: Array<String>) {
126131
" - Catalog: $catalogPath\n" +
127132
" - Release info: $releaseInfoPath\n" +
128133
" - Lexical DB: $lexicalDbPath\n" +
134+
" - Embed model: $embedModelPath\n" +
135+
" - Embed tokenizer: $embedTokenizerPath\n" +
129136
" - Text index: $textIndexDir\n" +
130137
" - Lookup index: $lookupIndexDir\n" +
131138
" -> Bundle .tar.zst: $bundleOutputPath\n" +
@@ -171,14 +178,28 @@ fun main(args: Array<String>) {
171178
logger.w { "Lexical DB missing: $lexicalDbPath (skipped)" }
172179
}
173180

181+
// Add the dense embedding model + tokenizer if available
182+
if (embedModelPath.exists()) {
183+
addFileToTar(tar, embedModelPath, embedModelPath.fileName.toString(), logger)
184+
logger.i { "Added embedding model to bundle" }
185+
} else {
186+
logger.w { "Embedding model missing: $embedModelPath (skipped, dense search disabled)" }
187+
}
188+
if (embedTokenizerPath.exists()) {
189+
addFileToTar(tar, embedTokenizerPath, embedTokenizerPath.fileName.toString(), logger)
190+
logger.i { "Added embedding tokenizer to bundle" }
191+
} else {
192+
logger.w { "Embedding tokenizer missing: $embedTokenizerPath (skipped)" }
193+
}
194+
174195
// Add the precomputed catalog if available
175196
if (haveCatalog) {
176197
addFileToTar(tar, catalogPath, catalogPath.fileName.toString(), logger)
177198
logger.i { "Added precomputed catalog to bundle" }
178199
} else {
179200
logger.w { "Precomputed catalog missing: $catalogPath (skipped)" }
180201
}
181-
202+
182203
// Add the release info file if available
183204
if (haveReleaseInfo) {
184205
addFileToTar(tar, releaseInfoPath, releaseInfoPath.fileName.toString(), logger)

‎generator/searchindex/build.gradle.kts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,12 @@ tasks.register<JavaExec>("buildLuceneIndexDefault") {
6464
systemProperty("inMemoryDb", "true")
6565
}
6666

67+
// Optional: -PvectorsBin=/path (dir with ids.i64+vecs.f32+meta.txt) → SINGLE
68+
// fused index (text + dense KnnFloatVectorField per line).
69+
(project.findProperty("vectorsBin") as String?)?.let { systemProperty("vectorsBin", it) }
70+
// Optional: -PindexThreads=N to cap concurrent indexing threads (lower = less RAM).
71+
(project.findProperty("indexThreads") as String?)?.let { systemProperty("indexThreads", it) }
72+
6773
jvmArgs = listOf(
6874
"-Xmx$generatorHeap",
6975
"-XX:+UseG1GC",
@@ -73,3 +79,23 @@ tasks.register<JavaExec>("buildLuceneIndexDefault") {
7379
)
7480
}
7581

82+
83+
// Build the DENSE vector Lucene index from SeforimEmbedding binaries (embed_corpus_bin.py).
84+
// ./gradlew :SeforimLibrary:searchindex:buildVectorIndex \
85+
// -DvecBinDir=/path/vector_index -DvecIndexOut=/path/vector-lucene
86+
tasks.register<JavaExec>("buildVectorIndex") {
87+
group = "application"
88+
description = "Build the dense vector Lucene index from embedding binaries."
89+
dependsOn("jvmJar")
90+
mainClass.set("io.github.kdroidfilter.seforimlibrary.searchindex.BuildVectorIndexKt")
91+
classpath = files(tasks.named("jvmJar")) + configurations.getByName("jvmRuntimeClasspath")
92+
// Big heap + G1 + Panama Vector SIMD (jdk.incubator.vector) → fast parallel HNSW build.
93+
jvmArgs = listOf(
94+
"-Xmx12g",
95+
"-XX:+UseG1GC",
96+
"--enable-native-access=ALL-UNNAMED",
97+
"--add-modules=jdk.incubator.vector",
98+
)
99+
System.getProperty("vecBinDir")?.let { systemProperty("vecBinDir", it) }
100+
System.getProperty("vecIndexOut")?.let { systemProperty("vecIndexOut", it) }
101+
}

‎generator/searchindex/src/jvmMain/kotlin/io/github/kdroidfilter/seforimlibrary/searchindex/BuildLuceneIndex.kt‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,10 @@ import org.apache.lucene.analysis.miscellaneous.PerFieldAnalyzerWrapper
2222
import org.apache.lucene.analysis.ngram.NGramTokenFilter
2323
import org.jsoup.Jsoup
2424
import org.jsoup.safety.Safelist
25+
import java.io.DataInputStream
2526
import java.io.File
27+
import java.nio.ByteBuffer
28+
import java.nio.ByteOrder
2629
import java.nio.file.Files
2730
import java.nio.file.Path
2831
import java.nio.file.Paths
@@ -58,6 +61,12 @@ fun main() = runBlocking {
5861
runCatching { Files.createDirectories(indexDir) }
5962
runCatching { Files.createDirectories(lookupDir) }
6063

64+
// Optional dense embeddings -> SINGLE fused index. -DvectorsBin points to a dir
65+
// with ids.i64 + vecs.f32 + meta.txt (produced by embed_corpus_bin.py against this
66+
// same DB). Each line then also gets a KnnFloatVectorField in the text index.
67+
val vectorProvider: ((Long) -> FloatArray?)? =
68+
System.getProperty("vectorsBin")?.let { loadVectorProvider(Paths.get(it), logger) }
69+
6170
// Open repository (prefer in-memory for faster reads)
6271
val useMemoryDb = (System.getProperty("inMemoryDb") ?: "true") != "false"
6372
// Use a shared in-memory DB so multiple connections can read concurrently when multithreading
@@ -115,7 +124,7 @@ fun main() = runBlocking {
115124
)
116125
)
117126

118-
LuceneTextIndexWriter(indexDir, analyzer = analyzer).use { writer ->
127+
LuceneTextIndexWriter(indexDir, analyzer = analyzer, vectorProvider = vectorProvider).use { writer ->
119128
LuceneLookupIndexWriter(lookupDir, analyzer = analyzer).use { lookup ->
120129
val books = repo.getAllBooks()
121130
val indexThreads = (System.getProperty("indexThreads") ?: Runtime.getRuntime().availableProcessors().toString()).toInt().coerceAtLeast(1)
@@ -300,3 +309,29 @@ private fun sanitizeAcronymTerm(raw: String): String {
300309
if (raw.isEmpty()) return ""
301310
return normalizePostHtmlForIndex(raw)
302311
}
312+
313+
/**
314+
* Loads dense embeddings produced by embed_corpus_bin.py (ids.i64 + vecs.f32 + meta.txt,
315+
* little-endian) into RAM and returns a thread-safe lineId -> vector lookup, used by
316+
* LuceneTextIndexWriter to attach a KnnFloatVectorField per line (single fused index).
317+
*/
318+
private fun loadVectorProvider(dir: Path, logger: co.touchlab.kermit.Logger): ((Long) -> FloatArray?) {
319+
val meta = File(dir.toFile(), "meta.txt").readText().trim().split(" ")
320+
val n = meta[0].toInt()
321+
val dim = meta[1].toInt()
322+
logger.i { "Loading $n dense vectors (dim $dim) from $dir for the fused index" }
323+
val idsBuf = ByteBuffer.wrap(File(dir.toFile(), "ids.i64").readBytes()).order(ByteOrder.LITTLE_ENDIAN)
324+
val rowOf = HashMap<Long, Int>(n * 2)
325+
for (i in 0 until n) rowOf[idsBuf.long] = i
326+
val vecs = FloatArray(n * dim)
327+
DataInputStream(File(dir.toFile(), "vecs.f32").inputStream().buffered(1 shl 20)).use { din ->
328+
val rec = ByteArray(dim * 4)
329+
val bb = ByteBuffer.wrap(rec).order(ByteOrder.LITTLE_ENDIAN)
330+
for (i in 0 until n) {
331+
din.readFully(rec); bb.rewind()
332+
for (j in 0 until dim) vecs[i * dim + j] = bb.float
333+
}
334+
}
335+
logger.i { "Dense vectors loaded (${(n.toLong() * dim * 4L) / 1_000_000} MB in RAM)" }
336+
return { lineId -> rowOf[lineId]?.let { row -> vecs.copyOfRange(row * dim, row * dim + dim) } }
337+
}

0 commit comments

Comments
 (0)