digest: cache chunk embeddings; long documents answer from BM25 while embedding in the background
Cached chunk embeddings change what a long document answers with: BM25 first, embedding uncached chunks (or the whole document) in the background so the next question on the same doc scores hybrid. selectText now returns its mode and the number of chunks that had a vector at answer-build time; the digest cache stores only hybrid answers, never a provisional bm25 one. The background worker batches through a Dixie gate and WaitBackground polls for true idle. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
This commit is contained in:
@@ -0,0 +1,174 @@
|
||||
package digest
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// maxBgQueue is the number of documents the background worker may hold at once.
|
||||
const maxBgQueue = 4
|
||||
|
||||
// bgWorker is the single background goroutine that embeds the chunks of a long document after a
|
||||
// bm25 answer queued it. It owns a bounded, deduplicated queue and yields the Dixie gate to
|
||||
// interactive digests one batch at a time.
|
||||
type bgWorker struct {
|
||||
engine *Engine
|
||||
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
queue []embedJob
|
||||
seen map[string]struct{}
|
||||
inProg bool
|
||||
started bool
|
||||
}
|
||||
|
||||
type embedJob struct {
|
||||
chunks []string
|
||||
contentHash string
|
||||
callerClass string
|
||||
}
|
||||
|
||||
// worker returns the engine's background worker, creating it once. It does not start the goroutine.
|
||||
func (e *Engine) worker() *bgWorker {
|
||||
e.bgOnce.Do(func() {
|
||||
w := &bgWorker{engine: e}
|
||||
w.cond = sync.NewCond(&w.mu)
|
||||
w.seen = map[string]struct{}{}
|
||||
e.bg = w
|
||||
})
|
||||
return e.bg
|
||||
}
|
||||
|
||||
// queueForEmbed asks the background worker to embed a document's chunks. The document is
|
||||
// deduplicated by contentHash+"|"+callerClass; a full queue drops it (and logs) so a later
|
||||
// question re-queues it.
|
||||
func (e *Engine) queueForEmbed(chunks []string, contentHash, callerClass string) {
|
||||
if e == nil {
|
||||
return
|
||||
}
|
||||
w := e.worker()
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if !w.started {
|
||||
w.started = true
|
||||
go w.run()
|
||||
}
|
||||
key := contentHash + "|" + callerClass
|
||||
if _, ok := w.seen[key]; ok {
|
||||
return
|
||||
}
|
||||
if len(w.queue) >= maxBgQueue {
|
||||
// Leave it un-seen: a later question re-queues it.
|
||||
if e.Log != nil {
|
||||
e.Log.Warn("background embed queue full")
|
||||
}
|
||||
return
|
||||
}
|
||||
w.seen[key] = struct{}{}
|
||||
w.queue = append(w.queue, embedJob{chunks: chunks, contentHash: contentHash, callerClass: callerClass})
|
||||
w.cond.Signal()
|
||||
}
|
||||
|
||||
func (w *bgWorker) run() {
|
||||
for {
|
||||
w.mu.Lock()
|
||||
for len(w.queue) == 0 {
|
||||
w.cond.Wait()
|
||||
}
|
||||
job := w.queue[0]
|
||||
w.queue = w.queue[1:]
|
||||
w.inProg = true
|
||||
w.mu.Unlock()
|
||||
|
||||
w.process(job)
|
||||
|
||||
w.mu.Lock()
|
||||
w.inProg = false
|
||||
w.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// process embeds the first MaxDocChunks chunks that lack a vector, EmbedBatch at a time, yielding
|
||||
// the gate between batches and storing each vector. On any embedding error it drops the rest of
|
||||
// the document.
|
||||
func (w *bgWorker) process(job embedJob) {
|
||||
e := w.engine
|
||||
key := func(i int) string {
|
||||
return dcache.EmbedKey(job.contentHash+"/sp1200", i, e.embedModel(), job.callerClass)
|
||||
}
|
||||
var idx []int
|
||||
for i := range job.chunks {
|
||||
if i >= MaxDocChunks {
|
||||
break
|
||||
}
|
||||
if _, ok := e.Cache.GetEmbed(key(i)); ok {
|
||||
continue
|
||||
}
|
||||
idx = append(idx, i)
|
||||
if len(idx) >= MaxDocChunks {
|
||||
break
|
||||
}
|
||||
}
|
||||
for b := 0; b < len(idx); b += EmbedBatch {
|
||||
j := b + EmbedBatch
|
||||
if j > len(idx) {
|
||||
j = len(idx)
|
||||
}
|
||||
batch := make([]string, j-b)
|
||||
for k := range batch {
|
||||
batch[k] = job.chunks[idx[b+k]]
|
||||
}
|
||||
if !w.acquireGate() {
|
||||
if e.Log != nil {
|
||||
e.Log.Warn("background embed gate")
|
||||
}
|
||||
return
|
||||
}
|
||||
ectx, cancel := context.WithTimeout(context.Background(), e.timeout())
|
||||
vecs, err := e.embed(ectx, batch)
|
||||
cancel()
|
||||
e.Gate.Release()
|
||||
if err != nil {
|
||||
if e.Log != nil {
|
||||
e.Log.Warn("background embed", "err", truncateErr(err, 512))
|
||||
}
|
||||
return
|
||||
}
|
||||
for k := range batch {
|
||||
e.Cache.PutEmbed(key(idx[b+k]), vecs[k])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// acquireGate takes the Dixie gate, retrying once per second on a refusal so interactive digests
|
||||
// get turns. It returns false only if the gate cannot be taken after the first retry, so the
|
||||
// document is not stuck behind a permanently busy gate.
|
||||
func (w *bgWorker) acquireGate() bool {
|
||||
ok, _ := w.engine.Gate.Acquire(context.Background())
|
||||
if ok {
|
||||
return true
|
||||
}
|
||||
time.Sleep(1 * time.Second)
|
||||
ok, _ = w.engine.Gate.Acquire(context.Background())
|
||||
return ok
|
||||
}
|
||||
|
||||
// WaitBackground returns when the background worker is idle: the queue empty and nothing embedding.
|
||||
func (e *Engine) WaitBackground(ctx context.Context) error {
|
||||
w := e.worker()
|
||||
for {
|
||||
w.mu.Lock()
|
||||
idle := len(w.queue) == 0 && !w.inProg
|
||||
w.mu.Unlock()
|
||||
if idle {
|
||||
return nil
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user