diff --git a/README.md b/README.md index 219c8043..05db03ee 100644 --- a/README.md +++ b/README.md @@ -205,11 +205,12 @@ carry their own PDF, video and image readers. **Will it slow my computer down?** Only the first pass is heavy. After that it reads new files only. The model runs in a separate process that quits when it's done, and the search side only loads while you're searching. -**Which languages?** The model knows more than 100, and DigUp searches in all of them, across languages too. The evals -cover only English, Bengali and Arabic because those are the languages of the files DigUp was tested on, and Google -notes that the model isn't equally strong in every language. Two limits apply to exact words, not to meaning: the -text inside screenshots is read only on a Mac, by Apple's text recognition, which knows about 25 languages, and -languages written without spaces, like Chinese, Japanese and Thai, aren't split into words. +**Which languages?** The model knows more than 100, and DigUp searches in all of them, across languages too. The semantic +evals cover only English, Bengali and Arabic because those are the languages of the files DigUp was tested on, and +Google notes that the model isn't equally strong in every language. Exact-word search reads screenshot text only on +a Mac, using Apple's text recognition (about 25 languages). Han characters in text bodies match as contiguous +sequences; other text written without spaces, such as Japanese kana and Thai, is not split into words. See +[Han body search](docs/han-body-search.md) for index-upgrade details. **How do I remove it?** On a Mac, quit it, then delete DigUp.app and `~/Library/Application Support/DigUp`, which holds the model and the diff --git a/Sources/DigUpKit/CodeExtractor.swift b/Sources/DigUpKit/CodeExtractor.swift index 0dfc0010..80762a96 100644 --- a/Sources/DigUpKit/CodeExtractor.swift +++ b/Sources/DigUpKit/CodeExtractor.swift @@ -23,9 +23,6 @@ public enum CodeExtractor { let cells: Int? } - /// The model reads 2,048 tokens of an input at most; a single line longer than this is cut here. - static let longestPiece = 4000 - static func read(_ url: URL, options: IndexOptions) throws -> File { let data = try Data(contentsOf: url) let (repo, title) = place(of: url.path) @@ -117,8 +114,8 @@ public enum CodeExtractor { } var last = end while last > start + 1, blank(last - 1) { last -= 1 } // blank lines at its end don't count - var body = lines[start.. longestPiece { body = String(body.prefix(longestPiece)) } + // The indexer splits complete prompts using the runtime tokenizer. Never discard a long line here. + let body = lines[start.. [SQLValue] { - let normalized = SearchText.normalized(text) + let normalized = SearchText.indexed(text) return [.text(normalized), normalized == text ? .null : .text(text)] } private func addNameRow(file: Int64, path: String) throws { try db.run("INSERT INTO segments(file, kind, modality) VALUES(?, 'name', 'none')", [.int(file)]) - try db.run("INSERT INTO fts(rowid, name, body) VALUES(?, ?, ?)", - [.int(db.lastInsertID), .text(SearchText.normalized(Self.nameText(path))), - .text(SearchText.normalized(Self.folderText(path)))]) + try db.run("INSERT INTO fts(rowid, name, body, original) VALUES(?, ?, ?, ?)", + [.int(db.lastInsertID), .text(SearchText.normalized(Self.nameText(path)))] + + Self.keywordText(Self.folderText(path))) } private func deleteSegments(of file: Int64, keepingName: Bool) throws { diff --git a/Sources/DigUpKit/Indexer.swift b/Sources/DigUpKit/Indexer.swift index 8c5bb405..3a83197b 100644 --- a/Sources/DigUpKit/Indexer.swift +++ b/Sources/DigUpKit/Indexer.swift @@ -583,28 +583,41 @@ public final class Indexer { let wordPages = textPages.filter { SearchText.hasWords($0.text) } let imagePages = pages.filter { $0.imagePath != nil } let garbledPages = pages.filter { $0.garbled != nil } // rendered too: among `imagePages` - var textVectors: [Int: [Float]] = [:] // by page number + var textParts: [(page: Int, text: String)] = [] + var textVectors: [Int: [Float]] = [:] // by part, not page (a page can have several) var imageVectors: [Embedding] = [] + var garbledParts: [(page: Int, text: String)] = [] var garbledVectors: [[Float]] = [] if !wordPages.isEmpty || !imagePages.isEmpty { let worker = try fullWorker() - let vectors = try worker.embedTexts(wordPages.map { - Prompts.document(title: title, text: String($0.text.prefix(4000))) + for page in textPages { + let parts = SearchText.hasWords(page.text) + ? try worker.documentChunks(page.text, title: title) : [page.text] + textParts += parts.map { (page.number, $0) } + } + let worded = textParts.indices.filter { SearchText.hasWords(textParts[$0].text) } + let vectors = try worker.embedTexts(worded.map { + Prompts.document(title: title, text: textParts[$0].text) }) - textVectors = Dictionary(uniqueKeysWithValues: zip(wordPages.map(\.number), vectors)) + textVectors = Dictionary(uniqueKeysWithValues: zip(worded, vectors)) imageVectors = try worker.embedImages(imagePages.compactMap(\.imagePath), budget: options.pageBudget) - garbledVectors = try worker.embedTexts(garbledPages.map { - Prompts.document(title: title, text: String($0.garbled!.prefix(4000))) + for page in garbledPages { + garbledParts += try worker.documentChunks(page.garbled!, title: title).map { (page.number, $0) } + } + garbledVectors = try worker.embedTexts(garbledParts.map { + Prompts.document(title: title, text: $0.text) }) + } else { + textParts = textPages.map { ($0.number, $0.text) } } try store.db.transaction { // Meaning only: the garbled words are never shown or matched against what you type. - for (page, vector) in zip(garbledPages, garbledVectors) { - try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(page.number), vector: vector) + for (part, vector) in zip(garbledParts, garbledVectors) { + try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(part.page), vector: vector) } - for page in textPages { - try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(page.number), - excerpt: excerpt(page.text), text: page.text, vector: textVectors[page.number]) + for (index, part) in textParts.enumerated() { + try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(part.page), + excerpt: excerpt(part.text), text: part.text, vector: textVectors[index]) } for (page, embedding) in zip(imagePages, imageVectors) { try store.addSegment(file: file.id, kind: .pageImage, modality: .image, loc: Double(page.number), @@ -634,18 +647,24 @@ public final class Indexer { } let title = URL(fileURLWithPath: file.path).deletingPathExtension().lastPathComponent // Text that's mostly digits (a data dump saved as .txt) gets keyword search only. - let worded = chunks.indices.filter { SearchText.hasWords(chunks[$0]) } - let vectors = worded.isEmpty ? [] : try fullWorker().embedTexts(worded.map { - Prompts.document(title: title, text: chunks[$0]) + let worker = chunks.contains(where: SearchText.hasWords) ? try fullWorker() : nil + var parts: [(offset: Int, text: String)] = [] + for (offset, chunk) in chunks.enumerated() { + let split = SearchText.hasWords(chunk) ? try worker!.documentChunks(chunk, title: title) : [chunk] + parts += split.map { (offset, $0) } + } + let worded = parts.indices.filter { SearchText.hasWords(parts[$0].text) } + let vectors = worded.isEmpty ? [] : try worker!.embedTexts(worded.map { + Prompts.document(title: title, text: parts[$0].text) }) let vector = Dictionary(uniqueKeysWithValues: zip(worded, vectors)) let end = min(total, options.maxChunks) let unread = read.upperBound < end try store.db.transaction { - for (offset, chunk) in chunks.enumerated() { + for (index, part) in parts.enumerated() { try store.addSegment(file: file.id, kind: .chunk, modality: .text, - loc: Double(read.lowerBound + offset + 1), excerpt: excerpt(chunk), text: chunk, - vector: vector[offset]) + loc: Double(read.lowerBound + part.offset + 1), excerpt: excerpt(part.text), text: part.text, + vector: vector[index]) } var info = ["chunks": "\(end)", "chars": "\(characters)"] info["read_to"] = unread ? "\(read.upperBound)" : nil @@ -656,7 +675,7 @@ public final class Indexer { let span = later ? "chunks \(read.lowerBound + 1)–\(read.upperBound) of \(end)" : unread ? "\(chunks.count) of \(end) chunks" : "\(chunks.count) chunks" return (unread ? "partial" : "done", - span + (worded.count < chunks.count ? ", \(chunks.count - worded.count) for keywords only" : "")) + span + (worded.count < parts.count ? ", \(parts.count - worded.count) for keywords only" : "")) } /// A code file's stretches of lines (a notebook's of cells), each embedded with the file's path in its repo as its @@ -666,14 +685,39 @@ public final class Indexer { try store.finish(file.id, state: "skipped", detail: "no code") return ("skipped", "no code") } - // Lines that are mostly numbers and symbols (a table of constants) get keyword search only. - let worded = code.pieces.indices.filter { SearchText.hasWords(code.pieces[$0].text) } - let vectors = worded.isEmpty ? [] : try fullWorker().embedTexts(worded.map { - Prompts.document(title: code.title, text: code.pieces[$0].text) + // Keep the extractor's coarse code boundaries, then apply the complete + // prompt's real token budget. Preserve source offsets through this split. + let worker = code.pieces.contains { SearchText.hasWords($0.text) } ? try fullWorker() : nil + var pieces: [CodeExtractor.Piece] = [] + for piece in code.pieces { + guard SearchText.hasWords(piece.text), let worker else { + pieces.append(piece) + continue + } + let ranges = try TextBudget.ranges(piece.text, title: code.title, limit: worker.textTokenLimit, + count: worker.textTokenCount) + var countedThrough = piece.text.startIndex, lineOffset = 0 + for range in ranges { + let text = String(piece.text[range]) + var first = piece.first, last = piece.last + if code.cells == nil { + // Ranges advance even with overlap; count each source prefix only once. + lineOffset += piece.text[countedThrough.. (String, String?) { diff --git a/Sources/DigUpKit/SearchText.swift b/Sources/DigUpKit/SearchText.swift index 35f06bad..8b521b22 100644 --- a/Sources/DigUpKit/SearchText.swift +++ b/Sources/DigUpKit/SearchText.swift @@ -2,6 +2,36 @@ import Foundation /// Text the way keyword search keeps it and matches it, so that a word is found however it was typed or extracted. public enum SearchText { + /// Body text and body query phrases only; filenames retain their unspaced representation. Han characters + /// become adjacent tokens, so "租房" matches inside "租房合同" without losing two-character words. + public static func indexed(_ text: String) -> String { + var result = "" + let characters = Array(normalized(text)) + for (index, character) in characters.enumerated() { + if isHan(character) { result += " " + String(character) + " " } + else { + result.append(character) + let boundary = !character.unicodeScalars.contains { CharacterSet.alphanumerics.contains($0) } + let besideHan = (index > 0 && isHan(characters[index - 1])) + || (index + 1 < characters.count && isHan(characters[index + 1])) + // Co is indexed but is never a query word. Source separators must + // not disappear between adjacent Han tokens. + if boundary && besideHan { result += " \u{E000} " } + } + } + return result.replacingOccurrences(of: #"[ \t]+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespaces) + } + + public static func isHan(_ character: Character) -> Bool { + character.unicodeScalars.contains { + switch $0.value { + case 0x3400...0x4DBF, 0x4E00...0x9FFF, 0xF900...0xFAFF, 0x20000...0x2FFFF, 0x30000...0x323AF: true + default: false + } + } + } + /// NFKC (Arabic presentation forms and ligatures become plain letters, full-width becomes ASCII, and Bengali's /// two spellings of "ো" become one), invisible formatting characters dropped (Bengali's zero-width joiners, /// direction marks), and Arabic the way people type it: no tatweel, no short-vowel marks, one alef. @@ -59,13 +89,16 @@ public enum SearchText { guard hit.location != NSNotFound else { break } searchRange = NSRange(location: hit.upperBound, length: haystack.length - hit.upperBound) guard var range = folded.original(hit) else { continue } - if range.lowerBound > text.startIndex, isWordCharacter(text[text.index(before: range.lowerBound)]) { + if word.first.map(isHan) != true, + range.lowerBound > text.startIndex, isWordCharacter(text[text.index(before: range.lowerBound)]), + !isHan(text[text.index(before: range.lowerBound)]) { continue } - if range.upperBound < text.endIndex, isWordCharacter(text[range.upperBound]) { + if word.last.map(isHan) != true, + range.upperBound < text.endIndex, isWordCharacter(text[range.upperBound]), !isHan(text[range.upperBound]) { guard prefix else { continue } var end = range.upperBound - while end < text.endIndex, isWordCharacter(text[end]) { end = text.index(after: end) } + while end < text.endIndex, isWordCharacter(text[end]), !isHan(text[end]) { end = text.index(after: end) } range = range.lowerBound.. String? { - let marks = terms.flatMap { matches(of: $0, in: text) }.sorted { $0.lowerBound < $1.lowerBound } + let marks = terms.flatMap { matches(of: $0, in: text) } + .sorted { $0.lowerBound < $1.lowerBound } guard let first = marks.first else { return nil } // About `words` words, starting a few before the first match. - var start = first.lowerBound, before = 0 - while start > text.startIndex, before < words / 3 { + let han = text.contains(where: isHan) + var start = first.lowerBound, before = 0, characters = 0 + while start > text.startIndex, before < words / 3, !han || characters < words * 3 { start = text.index(before: start) + characters += 1 if text[start].isWhitespace, start < first.lowerBound, !text[text.index(after: start)].isWhitespace { before += 1 } } if start > text.startIndex || text[start].isWhitespace { start = text.index(after: start) } var end = first.upperBound, after = 0 - while end < text.endIndex, after < words - before { + characters = 0 + while end < text.endIndex, after < words - before, !han || characters < words * 6 { if text[end].isWhitespace { after += 1 } end = text.index(after: end) + characters += 1 } var out = start > text.startIndex ? "…" : "" var cursor = start @@ -148,7 +186,7 @@ public enum Landmark { guard page > 1, texts.count >= page else { return nil } let target = texts[page - 1] let before = texts.prefix(page - 1).map { SearchText.normalized($0).lowercased() } - let onPage = words.filter { $0.count >= 3 }.compactMap { word in + let onPage = words.filter { $0.count >= 3 || ($0.count >= 2 && $0.contains(where: SearchText.isHan)) }.compactMap { word in SearchText.matches(of: word, in: target).first.map { String(target[$0]) } } return onPage.first { word in !before.contains { $0.contains(SearchText.normalized(word).lowercased()) } } diff --git a/Sources/DigUpKit/Searcher.swift b/Sources/DigUpKit/Searcher.swift index 4d7b6def..cc0325c7 100644 --- a/Sources/DigUpKit/Searcher.swift +++ b/Sources/DigUpKit/Searcher.swift @@ -143,8 +143,15 @@ public final class Searcher { // as typed stays whole, so "cat" doesn't bring in "category"; so does a word of two letters, the start of too // many words to stand for one. if identifier == nil { - for index in terms.indices where terms[index].count >= 3 && !patterns[index].hasSuffix("*") { - if try !isIndexed(patterns[index]) { patterns[index] += "*" } + for index in terms.indices where terms[index].count >= 3 + && !(prefixLast && index == terms.count - 1) && !SearchText.inflects(terms[index]) { + if try !isIndexed(patterns[index]) { + if terms[index].contains(where: SearchText.isHan) { + // Names are still whole Han tokens: retain their schema-5 missing-word prefix fallback. + // Body phrases must remain contiguous exact sequences of Han characters. + patterns[index] = "(name : \"\(terms[index])\"* OR body : \"\(SearchText.indexed(terms[index]))\")" + } else { patterns[index] += "*" } + } } } if !patterns.isEmpty { @@ -157,7 +164,8 @@ public final class Searcher { let file = row.int(0) guard words.count < 100, allowed(file), seen.insert(file).inserted else { return } // Text that normalizing changed is quoted as it's written (with its hamza, harakat, joiners). - let snippet = row.text(3).flatMap { SearchText.snippet($0, terms: terms) } ?? row.text(2) + // Stemming may match without a literal highlight. Never show the indexed Han separators. + let snippet = row.text(3).map { SearchText.snippet($0, terms: terms) ?? excerpt($0) } ?? row.text(2) words.append((file, row.int(1), snippet)) } if !words.isEmpty { @@ -228,7 +236,7 @@ public final class Searcher { // never shows page 2's words. A segment with no place and no text of its own (a screenshot's image // vector) can borrow them from its file's OCR text, which describes the same picture. var snippet = keyword.flatMap { $0.1.segment == segmentID ? $0.1.snippet : nil } - .flatMap { $0.contains("«") ? excerpt($0, limit: 400) : nil } + .flatMap { $0.contains("«") || segment != .name ? excerpt($0, limit: 400) : nil } if snippet == nil, storedExcerpt == nil, loc == nil, let match = keyword?.1, let text = match.snippet, text.contains("«") { var matchedKind: String? @@ -394,7 +402,12 @@ public final class Searcher { static func patterns(_ terms: [String], prefixLast: Bool) -> [String] { terms.enumerated().map { index, term in let prefix = (prefixLast && index == terms.count - 1) || SearchText.inflects(term) - return "\"\(term)\"" + (prefix ? "*" : "") + let suffix = prefix ? "*" : "" + if term.contains(where: SearchText.isHan) { + // Keep name matching (and the filename substring fallback) compatible with schema 5. + return "(name : \"\(term)\"\(suffix) OR body : \"\(SearchText.indexed(term))\"\(suffix))" + } + return "\"\(term)\"" + suffix } } diff --git a/Sources/DigUpKit/TextBudget.swift b/Sources/DigUpKit/TextBudget.swift new file mode 100644 index 00000000..ddc98664 --- /dev/null +++ b/Sources/DigUpKit/TextBudget.swift @@ -0,0 +1,58 @@ +import Foundation + +/// Splits documents against a complete prompt's tokenizer budget. Callers provide +/// the runtime's real token count; character counts never decide what fits. +public enum TextBudget { + /// A single file cannot fit; unlike a failed worker, this must not stop the indexing queue. + public struct Failure: Error, CustomStringConvertible { + public let description: String + } + public static func chunks(_ text: String, title: String?, limit: Int, + count: (String) throws -> Int) throws -> [String] { + try ranges(text, title: title, limit: limit, count: count).map { String(text[$0]) } + } + + /// Original offsets let code retain source lines even when repeated text is split. + static func ranges(_ text: String, title: String?, limit: Int, + count: (String) throws -> Int) throws -> [Range] { + guard !text.isEmpty else { return [] } + func fits(_ body: String) throws -> Bool { + try count(Prompts.document(title: title, text: body)) <= limit + } + if try fits(text) { return [text.startIndex..] = [] + while start < indices.count - 1 { + // Bound tokenizer work even for enormous pages. This bound is only + // a probe size: every emitted piece is checked with the tokenizer. + var low = start + 1, high = min(indices.count - 1, start + max(1800, limit * 4)) + var end = start + while low <= high { + let probe = low + (high - low) / 2 + if try fits(String(text[indices[start].. start else { + throw Failure(description: "document title and one character exceed the \(limit)-token budget") + } + // Prefer a word boundary when one is nearby; no-space Han stays overlapped. + if end < indices.count - 1, + let space = text[indices[start].. (end - start) / 2 { + let candidate = start + text.distance(from: indices[start], to: space) + if try fits(String(text[indices[start].. [[Float]] func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] func embedAudio(_ paths: [String]) throws -> [Embedding] + var textTokenLimit: Int { get } + func textTokenCount(_ text: String) throws -> Int /// Frees the model. Safe to call twice. func close() } +public extension Embedder { + var textTokenLimit: Int { info.textTokenLimit ?? 2048 } + /// Conservative byte bound for simple/test embedders. Production runtimes + /// override this with their tokenizer, including its special tokens. + func textTokenCount(_ text: String) throws -> Int { text.utf8.count + 2 } + func documentChunks(_ text: String, title: String?) throws -> [String] { + try TextBudget.chunks(text, title: title, limit: textTokenLimit, count: textTokenCount) + } +} + /// How to start a worker process that speaks the JSON-lines protocol (see `EmbeddingWorker`). public struct WorkerConfig: Sendable { public var executable: URL @@ -131,6 +145,7 @@ public final class EmbeddingWorker: Embedder { let vectors: String? let errors: [String?]? let info: WorkerInfo? + let counts: [Int]? } public init(_ config: WorkerConfig, textOnly: Bool) throws { @@ -162,6 +177,23 @@ public final class EmbeddingWorker: Embedder { return try send("text", texts, budget: nil).map { $0.vector ?? [] } } + public func textTokenCount(_ text: String) throws -> Int { + let id = nextID + nextID += 1 + let line = try JSONEncoder().encode(Request(id: id, op: "token_count", inputs: [text], budget: nil)) + do { + try input.fileHandleForWriting.write(contentsOf: line + [0x0A]) + } catch { + throw WorkerError.died("can't write to the worker\(logTail())") + } + let reply = try readReply() + guard reply.id == id, reply.ok == true, reply.counts?.count == 1, + let count = reply.counts?.first, count >= 0 else { + throw WorkerError.failed(reply.error ?? "worker does not report tokenizer counts") + } + return count + } + public func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { guard !paths.isEmpty else { return [] } return try send("image", paths, budget: budget) diff --git a/Sources/LlamaRuntime/LlamaEmbedder.swift b/Sources/LlamaRuntime/LlamaEmbedder.swift index 952c78f1..0dcfcce0 100644 --- a/Sources/LlamaRuntime/LlamaEmbedder.swift +++ b/Sources/LlamaRuntime/LlamaEmbedder.swift @@ -58,7 +58,7 @@ public final class LlamaEmbedder: Embedder { throw WorkerError.failed("can't load \(models.text.path)") } // One input must fit one micro-batch: 2048 tokens (the reference's cut), which also holds a 1120-token image or - // a 750-token 30 s clip. Queries are short, so the query encoder reserves less and cuts at 512. + // a 750-token 30 s clip. Queries are short, so the query encoder reserves a smaller 512-token budget. let maxTokens = textOnly ? 512 : 2048 var params = llama_context_default_params() params.n_ctx = UInt32(maxTokens) @@ -104,7 +104,7 @@ public final class LlamaEmbedder: Embedder { self.maxTokens = maxTokens let known = Self.modelInfo info = WorkerInfo(model: known.model, revision: known.revision, runtime: known.runtime, dtype: known.dtype, - dim: embeddingSize, textOnly: textOnly, loadSeconds: nil) + dim: embeddingSize, textOnly: textOnly, loadSeconds: nil, textTokenLimit: maxTokens) // The first calls build the Metal pipelines: pay that now (the panel is opening). do { for _ in 0..<2 { _ = try embedTexts([Prompts.query("warm up")]) } @@ -114,7 +114,7 @@ public final class LlamaEmbedder: Embedder { } let seconds = (Date().timeIntervalSince(started) * 100).rounded() / 100 info = WorkerInfo(model: info.model, revision: info.revision, runtime: info.runtime, dtype: info.dtype, - dim: info.dim, textOnly: textOnly, loadSeconds: seconds) + dim: info.dim, textOnly: textOnly, loadSeconds: seconds, textTokenLimit: maxTokens) LlamaLog.shared.line("ready in \(seconds) s (\(textOnly ? "text-only" : "full"), llama.cpp \(Self.buildTag), " + "vectors \(Self.vectorVersion))") } @@ -136,6 +136,9 @@ public final class LlamaEmbedder: Embedder { var vectors: [[Float]] = [] for text in texts { let tokens = tokenize(text) + guard tokens.count <= maxTokens else { + throw WorkerError.failed("text has \(tokens.count) tokens; limit is \(maxTokens). Split it before embedding") + } if sequences == Self.maxSequences || self.tokens + tokens.count > maxTokens { vectors += try run() } begin() for token in tokens { try add(token: token) } @@ -144,6 +147,9 @@ public final class LlamaEmbedder: Embedder { return vectors } + public var textTokenLimit: Int { maxTokens } + public func textTokenCount(_ text: String) throws -> Int { tokenize(text).count } + public func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { guard media != nil else { throw WorkerError.failed("this embedder was started text-only") } guard Self.budgets.contains(budget) else { throw WorkerError.failed("budget must be one of \(Self.budgets)") } @@ -247,7 +253,7 @@ public final class LlamaEmbedder: Embedder { // MARK: Text model - /// ` text `, cut to `maxTokens` the way transformers cuts (keeping both special tokens). + /// ` text `, never silently truncated. Embedding validates the budget. private func tokenize(_ text: String) -> [llama_token] { let utf8 = Array(text.utf8CString) var tokens = [llama_token](repeating: 0, count: utf8.count + 8) @@ -257,9 +263,6 @@ public final class LlamaEmbedder: Embedder { count = llama_tokenize(vocab, utf8, Int32(utf8.count - 1), &tokens, Int32(tokens.count), true, true) } tokens.removeLast(tokens.count - Int(max(count, 0))) - if tokens.count > maxTokens { - tokens = Array(tokens[0..<(maxTokens - 1)]) + [llama_vocab_eos(vocab)] - } return tokens } diff --git a/Sources/digup/IndexService.swift b/Sources/digup/IndexService.swift index 5c5a7102..94e18be5 100644 --- a/Sources/digup/IndexService.swift +++ b/Sources/digup/IndexService.swift @@ -368,6 +368,8 @@ final class SharedEmbedder: @unchecked Sendable { let model: any Embedder init(_ model: any Embedder) { self.model = model } var info: WorkerInfo { model.info } + var textTokenLimit: Int { model.textTokenLimit } + func textTokenCount(_ text: String) throws -> Int { try model.textTokenCount(text) } func embedTexts(_ texts: [String]) throws -> [[Float]] { try model.embedTexts(texts) } func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { try model.embedImages(paths, budget: budget) diff --git a/Sources/digup/WorkerService.swift b/Sources/digup/WorkerService.swift index 47a7a8d1..10fe6594 100644 --- a/Sources/digup/WorkerService.swift +++ b/Sources/digup/WorkerService.swift @@ -17,6 +17,7 @@ func workerService(_ args: Arguments) -> Int32 { let infoFields: [String: Any] = [ "model": info.model, "revision": info.revision, "runtime": info.runtime, "dtype": info.dtype, "dim": info.dim, "text_only": info.textOnly, "load_seconds": info.loadSeconds ?? 0, + "text_token_limit": embedder.textTokenLimit, ] emit(["ready": true, "info": infoFields]) while let line = readLine() { @@ -36,6 +37,9 @@ func workerService(_ args: Arguments) -> Int32 { continue case "text": embeddings = try embedder.embedTexts(inputs).map { Embedding(vector: $0, error: nil) } + case "token_count": + emit(["id": id, "ok": true, "counts": try inputs.map(embedder.textTokenCount)]) + continue case "image": embeddings = try embedder.embedImages(inputs, budget: request["budget"] as? Int ?? 280) case "audio": diff --git a/Tests/DigUpKitTests/ChineseSearchTests.swift b/Tests/DigUpKitTests/ChineseSearchTests.swift new file mode 100644 index 00000000..8d25eb2d --- /dev/null +++ b/Tests/DigUpKitTests/ChineseSearchTests.swift @@ -0,0 +1,254 @@ +#if os(macOS) +import AppKit +#endif +import Foundation +import Testing +@testable import DigUpKit + +@Suite final class ChineseSearchTests { + let folder: URL = { + let caches = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask)[0] + let url = caches.appendingPathComponent("DigUpChineseTests-\(UUID().uuidString)") + try? FileManager.default.createDirectory(at: url, withIntermediateDirectories: true) + return url + }() + + deinit { try? FileManager.default.removeItem(at: folder) } + + @discardableResult + func add(_ store: IndexStore, path: String, body: String? = nil, inode: UInt64 = 1) throws -> Int64 { + let id = try store.insert(Candidate(path: path, kind: .doc, size: 1, mtime: 1, inode: inode, device: 1)) + if let body { try store.addSegment(file: id, kind: .chunk, modality: .text, loc: 1, text: body) } + return id + } + + @Test func chineseLexicalEvalOnPublicSyntheticCorpus() throws { + // Each case has its own index, so a keyword-only R@1 has one unambiguous answer. + let cases: [(query: String, name: String, body: String?)] = [ + ("租房", "notes.txt", "这是一份关于租房费用的公开合成记录"), + ("发票", "notes.txt", "本月发票金额已记录在公开合成表格中"), + ("预算", "notes.txt", "年度预算安排来自公开合成资料"), + ("北京", "notes.txt", "我们计划乘坐高铁前往北京游览"), + ("高铁", "notes.txt", "乘坐高铁前往南方旅行"), + ("苹果", "notes.txt", "桌上放着新鲜苹果和香蕉"), + ("𠮷野", "notes.txt", "公开合成记录中的𠮷野家店铺"), + ("年度租赁合同扫描备份", "年度租赁合同扫描备份.pdf", nil), + + ("payment error", "notes.txt", "A payment error appeared on the synthetic receipt."), + ] + var correct = 0 + for (index, item) in cases.enumerated() { + let store = try IndexStore(directory: folder.appendingPathComponent("eval-\(index)")) + let path = "/synthetic/\(item.name)" + try add(store, path: path, body: item.body) + let hits = try Searcher(store: store, loadVectors: false).search(item.query, queryVector: nil) + if hits.first?.path == path { correct += 1 } + #expect(hits.first?.path == path, "Keyword-only synthetic case: \(item.query)") + } + print("CHINESE_LEXICAL_EVAL keyword-only R@1: \(correct)/\(cases.count); no model inference") + } + + @Test func chinesePhrasesAreContiguousAndHighlightedAsWritten() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("phrases")) + try add(store, path: "/synthetic/good.txt", body: "租房合同已经签署", inode: 1) + try add(store, path: "/synthetic/bad.txt", body: "租赁住房合同", inode: 2) + let hits = try Searcher(store: store, loadVectors: false).search("租房", queryVector: nil) + #expect(hits.map(\.path) == ["/synthetic/good.txt"]) + #expect(hits.first?.excerpt?.contains("«租房»合同") == true) + #expect(Landmark.find(page: 2, texts: ["公开资料", "这份租房合同已签署"], words: ["合同"]) == "合同") + #expect(SearchText.matches(of: "day", in: "Sunday day").map { String("Sunday day"[$0]) } == ["day"]) + } + + @Test func latinTermsBesideHanHighlightOriginalText() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("mixed-scripts")) + try add(store, path: "/synthetic/mixed.txt", body: "学习Swift语言。") + let hit = try #require(try Searcher(store: store, loadVectors: false).search("swift", queryVector: nil).first) + #expect(hit.excerpt == "学习«Swift»语言。") + #expect(hit.excerpt?.contains("\u{E000}") == false) + } + + @Test func stemmingNeverShowsIndexedHanSeparators() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("stemmed-han")) + try add(store, path: "/synthetic/stemmed.txt", body: "中文 running。") + let hit = try #require(try Searcher(store: store, loadVectors: false).search("runs", queryVector: nil).first) + #expect(hit.excerpt == "中文 running。") + let text = String(repeating: "甲", count: 1000) + "Swift" + String(repeating: "乙", count: 1000) + let snippet = try #require(SearchText.snippet(text, terms: ["swift"])) + #expect(snippet.count < 150) + #expect(snippet.contains("«Swift»")) + } + + @Test func chineseChunksRetainOverlapWithoutWhitespace() { + let text = String(repeating: "甲", count: 499) + "租房合同" + String(repeating: "乙", count: 1100) + "尾部标记" + let chunks = DocExtractor.chunks(text, size: 500, overlap: 100, limit: 100) + #expect(chunks.count > 2) + #expect(chunks.allSatisfy { $0.count <= 500 }) + #expect(chunks[1].prefix(100) == chunks[0].suffix(100)) + #expect(chunks.contains { $0.contains("租房合同") }) // phrase spans the first hard boundary. + #expect(chunks.last?.hasSuffix("尾部标记") == true) + } + + @Test func hanPhrasesDoNotBridgeSourceSeparatorsOrChooseTheWrongPage() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("separators")) + let id = try add(store, path: "/synthetic/pages.pdf") + try store.addSegment(file: id, kind: .page, modality: .text, loc: 1, text: "租，房 租。房 租\n房") + let searcher = try Searcher(store: store, loadVectors: false) + #expect(try searcher.search("租房", queryVector: nil).isEmpty) + try store.addSegment(file: id, kind: .page, modality: .text, loc: 2, text: "租房合同已经签署") + let hits = try Searcher(store: store, loadVectors: false).search("租房", queryVector: nil) + #expect(hits.first?.loc == 2) + #expect(hits.first?.excerpt?.contains("«租房»") == true) + } + + @Test func schemaFiveKeepsStoredTextSearchableAfterUpgrade() throws { + let directory = folder.appendingPathComponent("migration") + var store: IndexStore? = try IndexStore(directory: directory) + let id = try add(store!, path: "/synthetic/年度合同备份.txt", body: "连续中文租房记录", inode: 1) + // Model the old unsplit word representation, retaining its original text. + try store!.db.run("UPDATE fts SET name = '年度合同备份.txt' WHERE rowid IN (SELECT id FROM segments WHERE kind = 'name')") + try store!.db.run("UPDATE fts SET body = '连续中文租房记录', original = NULL WHERE rowid IN (SELECT id FROM segments WHERE kind = 'chunk')") + try store!.finish(id, state: "done", detail: nil) + try store!.db.execute("PRAGMA user_version = 5") + store = nil + let reopened = try IndexStore(directory: directory) + let searcher = try Searcher(store: reopened, loadVectors: false) + #expect(try searcher.search("租房", queryVector: nil).map(\.path) == ["/synthetic/年度合同备份.txt"]) + #expect(try searcher.search("年度合同备份", queryVector: nil).map(\.path) == ["/synthetic/年度合同备份.txt"]) + } + + @Test func bodySegmentationLeavesFilenameSubstringFallbackCompatible() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("name-compatibility")) + let id = try add(store, path: "/synthetic/年度合同备份.txt", body: "连续中文租房记录") + func namesContaining(_ needle: String) throws -> [Int64] { + var files: [Int64] = [] + // Same unspaced representation required by the existing filename-fallback PR. + try store.db.query(""" + SELECT s.file FROM fts JOIN segments s ON s.id = fts.rowid + WHERE s.kind = 'name' AND instr(fts.name, ?) > 0 + """, [.text(needle)]) { files.append($0.int(0)) } + return files + } + #expect(try namesContaining("合同") == [id]) + #expect(try Searcher(store: store, loadVectors: false).search("年度合同", queryVector: nil).first?.path + == "/synthetic/年度合同备份.txt") + try store.move(id, to: "/synthetic/中文目录/年度发票备份.txt") + #expect(try namesContaining("合同").isEmpty) + #expect(try namesContaining("发票") == [id]) + let hits = try Searcher(store: store, loadVectors: false).search("目录", queryVector: nil) + #expect(hits.first?.path == "/synthetic/中文目录/年度发票备份.txt") + #expect(hits.first?.excerpt?.contains("«目录»") == true) + } + + @Test func schemaFivePreservesOfflineCachedTextAndVectors() throws { + let directory = folder.appendingPathComponent("migration-budget") + var store: IndexStore? = try IndexStore(directory: directory) + let short = try add(store!, path: "/synthetic/short.txt", body: "租房合同", inode: 1) + let long = try add(store!, path: "/synthetic/long.txt", body: String(repeating: "甲", count: 1800), inode: 2) + let code = try store!.insert(Candidate(path: "/synthetic/code.md", kind: .code, size: 9000, + mtime: 1, inode: 3, device: 1)) + try store!.addSegment(file: code, kind: .lines, modality: .text, loc: 1, locEnd: 1, + text: String(repeating: "a", count: 4000), vector: FakeEmbedder.vector(for: "old")) + for id in [short, long, code] { try store!.finish(id, state: "done") } + try store!.db.execute("PRAGMA user_version = 5") + store = nil + let reopened = try IndexStore(directory: directory) + var states: [Int64: String] = [:] + try reopened.db.query("SELECT id, state FROM files") { states[$0.int(0)] = $0.text(1) } + #expect(states[short] == "done") + #expect(states[long] == "done") + #expect(states[code] == "done") + var version: Int64 = 0 + try reopened.db.query("PRAGMA user_version") { version = $0.int(0) } + #expect(version == 6) + var savedVectors = 0 + try reopened.db.query("SELECT count(*) FROM segments WHERE file = ? AND vec IS NOT NULL", [.int(code)]) { + savedVectors = Int($0.int(0)) + } + #expect(savedVectors == 1) + var savedBody = "" + try reopened.db.query("SELECT coalesce(fts.original, fts.body) FROM fts JOIN segments s ON s.id = fts.rowid WHERE s.file = ? AND s.kind = 'chunk'", [.int(long)]) { + savedBody = $0.text(0) ?? "" + } + #expect(savedBody == String(repeating: "甲", count: 1800)) + #expect(try Searcher(store: reopened, loadVectors: false).search("租房", queryVector: nil).first?.path + == "/synthetic/short.txt") + } + + @Test func longDocumentTailReachesTheEmbedder() throws { + let files = folder.appendingPathComponent("document-files") + try FileManager.default.createDirectory(at: files, withIntermediateDirectories: true) + let url = files.appendingPathComponent("long.txt") + let text = String(repeating: "公开合成的中文资料", count: 800) + "尾部检索标记" + try text.write(to: url, atomically: true, encoding: .utf8) + let store = try IndexStore(directory: folder.appendingPathComponent("document-index")) + let embedder = FakeEmbedder() + var options = IndexOptions() + options.docStep = 1 + let indexer = try Indexer(store: store, options: options, makeEmbedder: { embedder }) + _ = try indexer.sync([files]) + #expect(try indexer.run()["partial"] == 1) + #expect(!embedder.textInputs.contains { $0.contains("尾部检索标记") }) + #expect(try indexer.finishLongFiles()["done"] == 1) + #expect(embedder.textInputs.contains { $0.contains("尾部检索标记") }) + #expect(embedder.textInputs.allSatisfy { $0.utf8.count + 2 <= embedder.textTokenLimit }) + #expect(try Searcher(store: store, loadVectors: false).search("尾部检索标记", queryVector: nil).first?.path == url.path) + } + + @Test func impossibleFileBudgetDoesNotBlockTheQueue() throws { + let files = folder.appendingPathComponent("budget-files") + try FileManager.default.createDirectory(at: files, withIntermediateDirectories: true) + for name in [String(repeating: "title", count: 20), "short"] { + try "Public text.".write(to: files.appendingPathComponent(name + ".txt"), atomically: true, encoding: .utf8) + } + let store = try IndexStore(directory: folder.appendingPathComponent("failed-budget-index")) + let embedder = SmallBudgetEmbedder() + let indexer = try Indexer(store: store, options: IndexOptions(), makeEmbedder: { embedder }) + _ = try indexer.sync([files]) + let result = try indexer.run() + #expect(result["failed"] == 1) + #expect(result["done"] == 1) + #expect(try store.pendingCount() == 0) + } + + #if os(macOS) + @Test func pdfTailIsEmbeddedAndKeepsTheOriginalPageNumber() throws { + let files = folder.appendingPathComponent("files") + try FileManager.default.createDirectory(at: files, withIntermediateDirectories: true) + let url = files.appendingPathComponent("synthetic-long-page.pdf") + var box = CGRect(x: 0, y: 0, width: 1600, height: 6000) + let context = try #require(CGContext(url as CFURL, mediaBox: &box, nil)) + let long = String(repeating: "这是一段公开合成的中文资料。", count: 600) + "尾部检索标记" + for text in [String(repeating: "First synthetic page. ", count: 20), long] { + context.beginPDFPage(nil) + NSGraphicsContext.saveGraphicsState() + NSGraphicsContext.current = NSGraphicsContext(cgContext: context, flipped: false) + (text as NSString).draw(in: CGRect(x: 60, y: 60, width: 1480, height: 5800), + withAttributes: [.font: NSFont.systemFont(ofSize: 12)]) + NSGraphicsContext.restoreGraphicsState() + context.endPDFPage() + } + context.closePDF() + let store = try IndexStore(directory: folder.appendingPathComponent("pdf-index")) + let embedder = FakeEmbedder() + let indexer = try Indexer(store: store, options: IndexOptions(), makeEmbedder: { embedder }) + _ = try indexer.sync([files]) + #expect(try indexer.run()["done"] == 1) + #expect(embedder.textInputs.contains { $0.contains("尾部检索标记") }) + let hits = try Searcher(store: store, loadVectors: false).search("标记", queryVector: nil) + #expect(hits.first?.path == url.path) + #expect(hits.first?.loc == 2) + } + + #endif + +} + +private final class SmallBudgetEmbedder: Embedder { + private let base = FakeEmbedder() + var info: WorkerInfo { base.info } + var textTokenLimit: Int { 80 } + func embedTexts(_ texts: [String]) throws -> [[Float]] { try base.embedTexts(texts) } + func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { try base.embedImages(paths, budget: budget) } + func embedAudio(_ paths: [String]) throws -> [Embedding] { try base.embedAudio(paths) } + func close() {} +} diff --git a/Tests/DigUpKitTests/CodeTests.swift b/Tests/DigUpKitTests/CodeTests.swift index f03cea13..4db23b96 100644 --- a/Tests/DigUpKitTests/CodeTests.swift +++ b/Tests/DigUpKitTests/CodeTests.swift @@ -72,10 +72,10 @@ import Testing #expect(pieces.last?.last == 102) } - @Test func aVeryLongLineIsCutForTheModel() { + @Test func aVeryLongLineIsPreservedForTokenizerSplitting() { let pieces = CodeExtractor.lines("let a = 1\n" + String(repeating: "y", count: 9000) + "\nlet b = 2\n", size: 1800) #expect(pieces.map(\.first) == [1, 2, 3]) - #expect(pieces[1].text.count == CodeExtractor.longestPiece) + #expect(pieces[1].text.count == 9000) } @Test func windowsLineEndingsAndBlankEdges() { @@ -228,6 +228,50 @@ import Testing // MARK: The code index + @Test func chineseCodePassagesRespectTokenBudgetAndKeepTheirLine() throws { + try write("chinese/.git/HEAD", "ref: refs/heads/main\n") + let longLine = String(repeating: "公开合成的中文资料", count: 600) + "尾部检索标记" + let url = try write("chinese/README.md", "# Public synthetic notes\n" + longLine + "\nFinal line.\n") + let store = try IndexStore(directory: folder.appendingPathComponent("budget-index"), file: IndexStore.codeFile) + let embedder = BudgetCheckingCodeEmbedder() + let indexer = try Indexer(store: store, options: codeOptions, makeEmbedder: { embedder }) + _ = try indexer.sync([folder.appendingPathComponent("chinese")]) + #expect(try indexer.run() == ["done": 1]) + #expect(embedder.inputs.count > 3) + #expect(embedder.inputs.allSatisfy { embedder.textTokenCount($0) <= embedder.textTokenLimit }) + #expect(embedder.inputs.contains { $0.contains("尾部检索标记") }) + let hits = try Searcher(store: store, loadVectors: false).search("尾部检索标记", queryVector: nil) + let hit = try #require(hits.first) + #expect(hit.path == url.path) + #expect(hit.segment == .lines && hit.loc == 2 && hit.locEnd == 2) + } + + @Test func chineseMultilineBudgetSplitKeepsOriginalFocusLine() throws { + try write("multiline/.git/HEAD", "ref: refs/heads/main\n") + // One character-sized code piece, but several tokenizer-sized parts. Repeated + // lines make searching for each part's text an ambiguous way to recover offsets. + let repeated = String(repeating: "公开合成的中文资料", count: 12) + let lines = ["# Public synthetic notes"] + Array(repeating: repeated, count: 12) + ["尾部检索标记"] + let text = lines.joined(separator: "\n") + "\n" + #expect(CodeExtractor.lines(text).count == 1) + let url = try write("multiline/README.md", text) + let store = try IndexStore(directory: folder.appendingPathComponent("multiline-index"), file: IndexStore.codeFile) + let embedder = BudgetCheckingCodeEmbedder() + let indexer = try Indexer(store: store, options: codeOptions, makeEmbedder: { embedder }) + _ = try indexer.sync([folder.appendingPathComponent("multiline")]) + #expect(try indexer.run() == ["done": 1]) + #expect(embedder.inputs.count > 1) + #expect(embedder.inputs.allSatisfy { embedder.textTokenCount($0) <= embedder.textTokenLimit }) + #expect(embedder.inputs.contains { $0.contains("尾部检索标记") }) + let hit = try #require(try Searcher(store: store, loadVectors: false) + .search("尾部检索标记", queryVector: nil).first) + #expect(hit.path == url.path && hit.segment == .lines) + #expect(hit.focus == 14) + let firstLine = try #require(hit.loc) + #expect(firstLine > 1) + #expect(hit.locEnd == 14) + } + @Test func codeIndexIsAFileOfItsOwn() throws { try write("billing/.git/HEAD", "ref: refs/heads/main\n") try write("billing/retry.py", [function("retryWithBackoff", body: 22), function("circuit", body: 22), @@ -339,3 +383,24 @@ import Testing #expect(estimate.bytes > 21 * 1600 && estimate.bytes < 200_000) } } + +/// Deterministic tokenizer stand-in that rejects overflow as both production runtimes do. +/// Kept local to code regressions so existing FakeEmbedder-based tests retain their behavior. +private final class BudgetCheckingCodeEmbedder: Embedder { + private let base = FakeEmbedder() + var info: WorkerInfo { base.info } + var textTokenLimit: Int { 2048 } + var inputs: [String] { base.textInputs } + func textTokenCount(_ text: String) -> Int { text.utf8.count + 2 } + func embedTexts(_ texts: [String]) throws -> [[Float]] { + guard texts.allSatisfy({ textTokenCount($0) <= textTokenLimit }) else { + throw WorkerError.failed("code regression: prompt exceeds tokenizer budget") + } + return try base.embedTexts(texts) + } + func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { + try base.embedImages(paths, budget: budget) + } + func embedAudio(_ paths: [String]) throws -> [Embedding] { try base.embedAudio(paths) } + func close() { base.close() } +} diff --git a/Tests/DigUpKitTests/FakeEmbedder.swift b/Tests/DigUpKitTests/FakeEmbedder.swift index 87771e17..43e80f90 100644 --- a/Tests/DigUpKitTests/FakeEmbedder.swift +++ b/Tests/DigUpKitTests/FakeEmbedder.swift @@ -6,6 +6,7 @@ final class FakeEmbedder: Embedder { let info = WorkerInfo(model: "fake", revision: "0000000", runtime: "test", dtype: "f32", dim: Vectors.dimension, textOnly: false, loadSeconds: nil) private(set) var texts = 0 + private(set) var textInputs: [String] = [] static func vector(for text: String) -> [Float] { var state = UInt64(5381) @@ -21,6 +22,7 @@ final class FakeEmbedder: Embedder { func embedTexts(_ texts: [String]) throws -> [[Float]] { self.texts += texts.count + textInputs += texts return texts.map(Self.vector(for:)) } diff --git a/Tests/DigUpKitTests/TextBudgetTests.swift b/Tests/DigUpKitTests/TextBudgetTests.swift new file mode 100644 index 00000000..83d873e2 --- /dev/null +++ b/Tests/DigUpKitTests/TextBudgetTests.swift @@ -0,0 +1,36 @@ +import Foundation +import Testing +@testable import DigUpKit + +@Suite struct TextBudgetTests { + @Test func budgetIncludesTitleAndKeepsOverlappingHanAndFinalContent() throws { + let text = String(repeating: "甲", count: 110) + "租房合同" + String(repeating: "乙", count: 300) + "尾部标记" + let title = "公开合成标题" + let count = { (input: String) in input.utf8.count + 2 } + let parts = try TextBudget.chunks(text, title: title, limit: 220, count: count) + #expect(parts.count > 3) + #expect(parts.allSatisfy { count(Prompts.document(title: title, text: $0)) <= 220 }) + #expect(parts.contains { $0.contains("租房合同") }) + #expect(parts.last?.hasSuffix("尾部标记") == true) + for pair in zip(parts, parts.dropFirst()) { + #expect((1...min(200, pair.0.count / 4)).contains { + pair.0.suffix($0) == pair.1.prefix($0) + }) + } + } + + @Test func smallEnglishAndBengaliArabicPassagesStayUnchanged() throws { + for text in ["A payment error appeared.", "সুন্দরবনে আমরা কুমিরও দেখেছি।", "قَالَ إِنَّ الْمَكْتَبَةَ تُغْلَقُ"] { + #expect(try TextBudget.chunks(text, title: "synthetic", limit: 2048, + count: { $0.utf8.count + 2 }) == [text]) + } + } + + @Test func impossibleTitleBudgetFailsInsteadOfDiscardingText() { + #expect(throws: (any Error).self) { + _ = try TextBudget.chunks("合同", title: String(repeating: "标题", count: 100), limit: 50, + count: { $0.utf8.count + 2 }) + } + } + +} diff --git a/Tests/DigUpKitTests/WorkerBudgetTests.swift b/Tests/DigUpKitTests/WorkerBudgetTests.swift new file mode 100644 index 00000000..c85d132a --- /dev/null +++ b/Tests/DigUpKitTests/WorkerBudgetTests.swift @@ -0,0 +1,56 @@ +#if !os(Windows) +import Foundation +import Testing +@testable import DigUpKit + +/// The JSON-lines contract, without a model or a Python environment. +@Suite struct WorkerBudgetTests { + private func withWorker(reply: String, limit: Int? = 77, _ test: (EmbeddingWorker) throws -> Void) throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: directory) } + let script = directory.appendingPathComponent("worker.sh") + let limitField = limit.map { ",\"text_token_limit\":\($0)" } ?? "" + let ready = "{\"ready\":true,\"info\":{\"model\":\"mock\",\"revision\":\"1\",\"runtime\":\"test\",\"dtype\":\"f32\",\"dim\":768,\"text_only\":true\(limitField)}}" + try """ + printf '%s\\n' '\(ready)' + while IFS= read -r line; do + case "$line" in + *token_count*) printf '%s\\n' '\(reply)' ;; + *quit*) exit 0 ;; + esac + done + """.write(to: script, atomically: true, encoding: .utf8) + let worker = try EmbeddingWorker(WorkerConfig(executable: URL(fileURLWithPath: "/bin/sh"), arguments: [script.path]), + textOnly: true) + defer { worker.close() } + try test(worker) + } + + @Test func countsAndSnakeCaseLimitRoundTrip() throws { + try withWorker(reply: #"{"id":1,"ok":true,"counts":[42]}"#) { worker in + #expect(worker.info.textOnly) + #expect(worker.info.textTokenLimit == 77) + #expect(worker.textTokenLimit == 77) + let count = try worker.textTokenCount("title: public | text: 租房") + #expect(count == 42) + } + } + + @Test func legacyUnsupportedCountFailsExplicitly() throws { + try withWorker(reply: #"{"id":1,"ok":false,"error":"unknown op token_count"}"#, limit: nil) { worker in + #expect(worker.textTokenLimit == 2048) + #expect(throws: WorkerError.self) { try worker.textTokenCount("public") } + } + } + + @Test(arguments: [#"{"id":2,"ok":true,"counts":[42]}"#, + #"{"id":1,"ok":true,"counts":[-1]}"#, + #"{"id":1,"ok":true,"counts":[1,2]}"#]) + func malformedCountsFailExplicitly(_ reply: String) throws { + try withWorker(reply: reply) { worker in + #expect(throws: WorkerError.self) { try worker.textTokenCount("public") } + } + } +} +#endif diff --git a/Tests/LlamaRuntimeTests/NativeTextBudgetTests.swift b/Tests/LlamaRuntimeTests/NativeTextBudgetTests.swift new file mode 100644 index 00000000..17d86845 --- /dev/null +++ b/Tests/LlamaRuntimeTests/NativeTextBudgetTests.swift @@ -0,0 +1,36 @@ +import DigUpKit +import Foundation +import Testing +@testable import LlamaRuntime + +@Suite(.serialized) struct NativeTextBudgetTests { + static let textModelPresent = (try? LlamaModels.locate()) != nil + + @Test(.enabled(if: textModelPresent)) func oversizedTextNeverSilentlyDropsItsTail() throws { + let embedder = try LlamaEmbedder(models: LlamaModels.locate(), textOnly: true) + defer { embedder.close() } + // Synthetic unspaced Han is far beyond the query encoder's 512-token context. + #expect(throws: (any Error).self) { + _ = try embedder.embedTexts([String(repeating: "租房合同", count: 1500) + "尾部标记"]) + } + } + + @Test(.enabled(if: textModelPresent)) func actualTokenizerBudgetPreservesEverySyntheticRecord() throws { + let embedder = try LlamaEmbedder(models: LlamaModels.locate(), textOnly: true) + defer { embedder.close() } + let records = (1...180).map { "公开记录\($0)的租房合同。" } + let parts = try embedder.documentChunks(records.joined() + "尾部标记", title: "公开合成回归") + #expect(parts.count > 1) + #expect(try parts.allSatisfy { + try embedder.textTokenCount(Prompts.document(title: "公开合成回归", text: $0)) <= embedder.textTokenLimit + }) + #expect(records.allSatisfy { record in parts.contains { $0.contains(record) } }) + #expect(parts.last?.hasSuffix("尾部标记") == true) + let vectors = try embedder.embedTexts(parts.prefix(2).map { + Prompts.document(title: "公开合成回归", text: $0) + }) + #expect(vectors.count == 2) + #expect(vectors.allSatisfy { $0.count == embedder.info.dim && $0.allSatisfy(\.isFinite) }) + print("NATIVE_TEXT_BUDGET actual tokenizer: \(parts.count) complete parts within \(embedder.textTokenLimit) tokens; 2 inference vectors validated") + } +} diff --git a/docs/han-body-search.md b/docs/han-body-search.md new file mode 100644 index 00000000..dbfcb113 --- /dev/null +++ b/docs/han-body-search.md @@ -0,0 +1,47 @@ +# Han body search and complete text embeddings + +This change is based on upstream `ec28973268b81f836be72187ff23b3a16ef4fd97`. + +## Scope + +- Index Han characters in document, PDF, OCR, code, and folder bodies as adjacent FTS tokens. Quoted query phrases preserve contiguity; punctuation and whitespace cannot disappear between characters. +- Keep `fts.name` normalized and **unspaced**, including after moves and migration. This preserves whole-name/prefix search and compatibility with the separate filename substring proposal (`instr(fts.name, ...)`). No filename trigram index is added. +- Highlight original text, including mixed Latin/Han text, without exposing indexing separators. Keep excerpts bounded for unspaced Han. +- Split complete title-plus-body prompts against the runtime tokenizer, including special tokens. Preserve document continuation offsets and PDF page/code line locations. Remove the old 4,000-character code-line cut and PDF text prefix cut. Both runtimes reject oversized requests instead of silently truncating them. +- Forward token counts through the JSON-lines worker and the indexing service's shared-model wrapper. + +No app identity, index location, media recovery/ranking, startup, localization, packaging, or model-version changes are included. + +## Existing indexes + +Schema 5 becomes schema 6 in one transaction. Bodies are rebuilt from `coalesce(original, body)`; names, vectors, file state, and segment locations remain intact. Migration does not read source files or require them to be connected. Older binaries cannot open schema 6, as with other schema upgrades. + +An old vector may already have lost a long passage's tail. Migration cannot reconstruct missing source text or replace those embeddings without reading the file and loading the tokenizer. The new splitting applies to new/changed files. To repair unchanged files, explicitly reindex when their folders are available, for example: + +```sh +digup index /path/to/folder --index /path/to/index --reindex +``` + +`--reindex` resets the selected index, so include all folders that belong in it; it is not a per-file repair command. Use the corresponding code index with `--code` when needed. Keep an index backup if rollback to an older binary is required. + +## Regression coverage + +`ChineseSearchTests`, `TextBudgetTests`, `CodeTests`, and `WorkerBudgetTests` cover body phrase contiguity, original snippets, supplementary-plane Han, filename compatibility, schema-5/offline-cache preservation, overlap, tail coverage, multi-pass documents, >4,000-character code lines, source line focus, impossible per-file budgets, and token-count protocol validation. + +`NativeTextBudgetTests` uses the actual pinned model when available. It checks oversized-input rejection, every synthetic record and final marker surviving a tokenizer-budget split, and finite 768-dimensional inference vectors. It explicitly skips without the model, which is not an inference pass. Python worker contract tests use deterministic stubs and do not load MLX: + +```sh +python -m unittest discover -s worker -p test_text_budget.py -v +``` + +### Measured before/after (Linux, Swift 6.2) + +- Same nine public synthetic keyword-only cases: upstream R@1 **2/9**, this change **9/9**. This is a narrow lexical regression corpus, not a multilingual semantic ranking benchmark. +- Actual pinned text model, query encoder's 512-token budget: upstream accepts a 3,013-token input and two different tails produce identical vectors. This change returns an explicit over-budget error. +- Actual tokenizer: 180 synthetic Han records plus a final marker become five complete parts within 512 tokens; every record survives and two inference vectors are validated. +- Full upstream suite: 92 tests, seven assertion issues. Full changed suite: 113 tests, the same seven baseline issues, no additional failures. Six are portable TIFF/HEIC/WebP orientation assertions with this environment's image reader; one is `placeInItsRepo` encountering the sandbox's synthetic `.git` parent. The final run excluding those two affected test functions passes all **111 selected tests with no skips**, including the actual-model budget tests and existing media/text `ReferenceVectorTests` with both pinned models. The long-document reference passes; no reference vectors or model version were changed. +- Python tokenizer contract: three tests pass. + +For Linux native model runs, set `DIGUP_MODELS` to the model folder and `DIGUP_BACKENDS` to `Vendor/llama/lib`; a development test binary does not carry ggml backend modules beside itself. All model files are pinned by `ModelFiles` and checksum-verified. + +The upstream `evals/` semantic suites need their generated `DigUpTestbed` source corpus, which is not bundled here. They are not claimed as run. macOS PDF extraction/rendering, Apple frameworks, and packaged-app validation remain separate from Linux engine validation. diff --git a/worker/digup_worker.py b/worker/digup_worker.py index 30158839..fa98293e 100644 --- a/worker/digup_worker.py +++ b/worker/digup_worker.py @@ -100,16 +100,25 @@ def __init__(self, model, text_only): "dim": DIM, "text_only": text_only, "load_seconds": round(time.perf_counter() - started, 2), + "text_token_limit": 512 if text_only else 2048, } def _vectors(self, output): return np.array(output.text_embeds.astype(self.mx.float32)) + def token_counts(self, inputs): + if not inputs: + return [] + tokens = self.tokenizer(inputs, padding=False, truncation=False, add_special_tokens=True) + return [len(ids) for ids in tokens["input_ids"]] + def text(self, inputs): + if any(count > self.info["text_token_limit"] for count in self.token_counts(inputs)): + raise ValueError("text exceeds the tokenizer budget; split it before embedding") rows = [] for i in range(0, len(inputs), BATCH["text"]): batch = inputs[i : i + BATCH["text"]] - tokens = self.tokenizer(batch, padding=True, truncation=True, max_length=2048, return_tensors="np") + tokens = self.tokenizer(batch, padding=True, truncation=False, return_tensors="np") rows.append(self._vectors(self.model(**{k: self.mx.array(v) for k, v in tokens.items()}))) return np.concatenate(rows), [None] * len(inputs) @@ -178,6 +187,9 @@ def main(): reply({"id": request_id, "ok": True, "info": embedder.info}) continue inputs = request.get("inputs") or [] + if op == "token_count": + reply({"id": request_id, "ok": True, "counts": embedder.token_counts(inputs)}) + continue if op == "text": vectors, errors = embedder.text(inputs) elif op in ("image", "audio"): diff --git a/worker/test_text_budget.py b/worker/test_text_budget.py new file mode 100644 index 00000000..5cb07e0d --- /dev/null +++ b/worker/test_text_budget.py @@ -0,0 +1,57 @@ +"""Tokenizer-budget contract without loading/downloading an MLX model.""" +import unittest +from types import SimpleNamespace + +import numpy as np + +from digup_worker import Embedder + + +class Tokenizer: + def __init__(self): + self.calls = [] + + def __call__(self, texts, **options): + self.calls.append(options) + rows = [[0] + list(text.encode('utf-8')) + [1] for text in texts] + if options.get('padding'): + width = max(map(len, rows)) + rows = [row + [0] * (width - len(row)) for row in rows] + return {'input_ids': rows} + + +class TextBudgetTests(unittest.TestCase): + def setUp(self): + self.embedder = Embedder.__new__(Embedder) + self.embedder.tokenizer = Tokenizer() + self.embedder.info = {'text_token_limit': 8} + self.embedder.mx = SimpleNamespace(array=np.array, float32=np.float32) + self.model_calls = [] + + def model(**tokens): + self.model_calls.append(tokens) + return SimpleNamespace(text_embeds=np.ones((len(tokens['input_ids']), 768), dtype=np.float32)) + + self.embedder.model = model + + def test_special_tokens_count_and_empty_batch(self): + self.assertEqual(self.embedder.token_counts(['abc', '租房']), [5, 8]) + self.assertEqual(self.embedder.token_counts([]), []) + self.assertTrue(self.embedder.tokenizer.calls[0]['add_special_tokens']) + self.assertFalse(self.embedder.tokenizer.calls[0]['truncation']) + + def test_exact_boundary_is_accepted_without_truncation(self): + vectors, errors = self.embedder.text(['租房']) + self.assertEqual(vectors.shape, (1, 768)) + self.assertEqual(errors, [None]) + self.assertEqual(self.model_calls[0]['input_ids'].shape, (1, 8)) + self.assertTrue(all(call['truncation'] is False for call in self.embedder.tokenizer.calls)) + + def test_overflow_rejected_before_any_model_call(self): + with self.assertRaisesRegex(ValueError, 'tokenizer budget'): + self.embedder.text(['ok', '租房x']) + self.assertEqual(self.model_calls, []) + + +if __name__ == '__main__': + unittest.main()