diff options
| -rw-r--r-- | llm/tokenbuffer.go | 26 |
1 files changed, 13 insertions, 13 deletions
diff --git a/llm/tokenbuffer.go b/llm/tokenbuffer.go index f83858d..19d7db9 100644 --- a/llm/tokenbuffer.go +++ b/llm/tokenbuffer.go @@ -43,14 +43,11 @@ func (tb *TokenBuffer) Seen() int { func (tb *TokenBuffer) Push(document int, text string) iter.Seq2[[]int64, int] { return func(yield func([]int64, int) bool) { if tb.document != -1 && document != tb.document { - if tb.includeTail && len(tb.buffer) > 0 { - if !yield(tb.buffer, tb.seen) { - return - } - } + tail, seen := tb.Tail() - tb.buffer = nil - tb.seen = 0 + if len(tail) > 0 && !yield(tail, seen) { + return + } } tb.document = document @@ -80,19 +77,22 @@ func (tb *TokenBuffer) Push(document int, text string) iter.Seq2[[]int64, int] { } func (tb *TokenBuffer) Tail() ([]int64, int) { + seen := tb.seen + + tb.document = -1 + tb.seen = 0 + if !tb.includeTail || len(tb.buffer) == 0 { - return nil, tb.seen + tb.buffer = tb.buffer[:0] + + return nil, seen } w := make([]int64, len(tb.buffer)) copy(w, tb.buffer) - seen := tb.seen - - tb.buffer = nil - tb.document = -1 - tb.seen = 0 + tb.buffer = tb.buffer[:0] return w, seen } |