1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
//go:build chunkdemo

// chunkdemo demonstrates why naive chunking of long posts loses meaning.
//
// The 256-token limit means a long post must either be truncated or split.
// Splitting looks harmless until you notice that a transformer embedding is
// contextual: every token's vector is computed in the presence of the other
// tokens in the same forward pass. Cut a text in half and the second half is
// embedded as though the first half never existed.
//
// This program makes that concrete by embedding the same words in different
// contexts and printing the cosine similarities, so the effect can be measured
// rather than argued about.
//
// It is a demo, not part of the program:
//
//	go run -tags chunkdemo .
//
// The build tag keeps it out of the ordinary build, so `nix build` and the
// eventual fedisearch binary never see it.
package main

import (
	"fmt"
	"math"
	"os"
	"strings"
)

func main() {
	if err := chunkdemo(); err != nil {
		fmt.Fprintf(os.Stderr, "chunkdemo: %v\n", err)
		os.Exit(1)
	}
}

func chunkdemo() error {
	modelDir, err := ensureModel()
	if err != nil {
		return err
	}
	emb, err := newEmbedder(modelDir)
	if err != nil {
		return err
	}
	defer emb.close()

	// embed returns unnormalized vectors; normalize so cosine is a dot product.
	embedUnit := func(s string) ([]float32, error) {
		v, err := emb.embed(s)
		if err != nil {
			return nil, err
		}
		return normalize(v), nil
	}

	fmt.Println("Each score is cosine similarity: 1.00 identical direction,")
	fmt.Println("0.00 unrelated. Same 384-dimensional space throughout.")

	// -----------------------------------------------------------------
	// 1. The premise: identical words, different neighbours.
	// -----------------------------------------------------------------
	//
	// If token vectors were fixed per word, these two sentences would embed
	// identically, because they are anagrams of each other at the word level.
	section("1. Word order changes meaning, and the model sees it")
	if err := compare(embedUnit,
		"dog bites man",
		"man bites dog",
	); err != nil {
		return err
	}
	fmt.Println("   Same three words. A bag-of-words model would score 1.00 here.")

	// -----------------------------------------------------------------
	// 2. The actual failure: a fragment cut loose from its context.
	// -----------------------------------------------------------------
	//
	// This is the user's example. The full sentence is about a *film*. The
	// second half, read alone, is about violence. If a chunker splits at the
	// wrong point, the second chunk is embedded with no idea a film was ever
	// mentioned.
	section("2. A fragment split from its context loses the subject")
	const full = "I went shopping for some movie that talked about splattering his brains out"
	const head = "I went shopping for some movie that"
	const tail = "talked about splattering his brains out"

	if err := compare(embedUnit, full, tail); err != nil {
		return err
	}
	fmt.Println("   The tail alone is a poor stand-in for the whole sentence.")

	// Does the tail still look like it is about films? Compare both against a
	// query someone might plausibly type.
	const filmQuery = "a film recommendation"
	fmt.Println()
	fmt.Printf("   Against the query %q:\n", filmQuery)
	if err := compareLabeled(embedUnit, filmQuery, full, "full sentence"); err != nil {
		return err
	}
	if err := compareLabeled(embedUnit, filmQuery, tail, "tail chunk only"); err != nil {
		return err
	}
	if err := compareLabeled(embedUnit, filmQuery, head, "head chunk only"); err != nil {
		return err
	}
	fmt.Println("   If the tail scores far lower, the film context was lost with the cut.")

	const violenceQuery = "graphic violence and gore"
	fmt.Println()
	fmt.Printf("   Against the query %q:\n", violenceQuery)
	if err := compareLabeled(embedUnit, violenceQuery, full, "full sentence"); err != nil {
		return err
	}
	if err := compareLabeled(embedUnit, violenceQuery, tail, "tail chunk only"); err != nil {
		return err
	}
	fmt.Println("   If the tail scores higher than the full sentence, the chunk is")
	fmt.Println("   actively misleading: it reads as real violence, not a film plot.")

	// -----------------------------------------------------------------
	// 3. Averaging chunk vectors does not reconstruct the whole.
	// -----------------------------------------------------------------
	//
	// A tempting fix is to embed each chunk and average the results. That is
	// not the same as embedding the whole text: the averaging happens after
	// the transformer, so no information ever flows between the chunks.
	section("3. Averaging chunk vectors is not the same as embedding the whole")
	vFull, err := embedUnit(full)
	if err != nil {
		return err
	}
	vHead, err := embedUnit(head)
	if err != nil {
		return err
	}
	vTail, err := embedUnit(tail)
	if err != nil {
		return err
	}
	avg := normalize(mean(vHead, vTail))
	fmt.Printf("   %-46s %.4f\n", "mean(head, tail)  vs  full sentence", dot(avg, vFull))
	fmt.Println("   Close, but not 1.00: the reconstruction is lossy, and the loss")
	fmt.Println("   grows with how much the halves depend on each other.")

	// -----------------------------------------------------------------
	// 4. Overlap is the standard mitigation.
	// -----------------------------------------------------------------
	//
	// Repeating a window of tokens at the chunk boundary gives the second
	// chunk some of the first chunk's context. It costs storage and still
	// does not recover long-range dependencies, but it is cheap.
	section("4. Overlapping the split recovers much of the context")
	const tailOverlap = "for some movie that talked about splattering his brains out"
	if err := compareLabeled(embedUnit, filmQuery, tail, "tail, no overlap"); err != nil {
		return err
	}
	if err := compareLabeled(embedUnit, filmQuery, tailOverlap, "tail, 4 words overlap"); err != nil {
		return err
	}
	fmt.Println("   Carrying a few words across the boundary restores the subject.")

	// -----------------------------------------------------------------
	// 5. Where truncation would actually bite.
	// -----------------------------------------------------------------
	//
	// Truncation is the other option, and it is not obviously worse: it keeps
	// the beginning, which for a post is usually where the topic is stated.
	section("5. Truncation keeps the topic when the topic comes first")
	const longish = "Today I want to talk about database indexes. " +
		"There is a lot of detail below but the short version is that " +
		"most people reach for one far too late."
	const firstOnly = "Today I want to talk about database indexes."
	if err := compareLabeled(embedUnit, "database indexing", longish, "whole post"); err != nil {
		return err
	}
	if err := compareLabeled(embedUnit, "database indexing", firstOnly, "first sentence only"); err != nil {
		return err
	}
	fmt.Println("   For posts that open with their subject, truncation is mild.")
	fmt.Println("   For posts that bury the point, it is not.")

	// -----------------------------------------------------------------
	// 6. Can the first chunk's vector be "injected" into the second?
	// -----------------------------------------------------------------
	//
	// The appealing idea: embed chunk 1, then somehow hand that vector to
	// chunk 2 as a prior, so the second forward pass knows what came before.
	//
	// The honest answer is that you cannot inject it into the *weights*
	// without retraining — the weights are what the model learned, and are
	// identical for every input. But there are three cheaper things that are
	// often confused with it, and they differ a lot in how well they work.
	// The distinction matters, so measure all of them.
	section("6. Carrying context across a split: what actually works")

	vQuery, err := embedUnit(filmQuery)
	if err != nil {
		return err
	}
	score := func(v []float32) float32 { return dot(vQuery, v) }

	// (a) Baseline: the tail alone, no context at all.
	fmt.Printf("     %.4f  (a) tail alone\n", score(vTail))

	// (b) Vector arithmetic: average the two chunk embeddings. This is what
	//     "injecting a prior" usually degrades into. It happens entirely
	//     after both forward passes, so no token in chunk 2 was ever computed
	//     in the presence of chunk 1. It shifts the result toward chunk 1
	//     without informing it.
	fmt.Printf("     %.4f  (b) mean(head, tail) — arithmetic after the fact\n", score(avg))

	// (c) Weighted blend, favouring the tail. Same objection as (b): the
	//     numbers move but nothing was recomputed. Included because it is the
	//     obvious next thing to try, and it is worth seeing that tuning the
	//     weight does not rescue the approach.
	blend := normalize(weightedMean(vHead, vTail, 0.3, 0.7))
	fmt.Printf("     %.4f  (c) 0.3*head + 0.7*tail — tuning the weight\n", score(blend))

	// (d) Text-level prefixing: prepend a short *summary* of the context and
	//     run a genuine forward pass. Unlike (b) and (c) this recomputes every
	//     token in the presence of the context, so "movie" can actually
	//     condition "splattering". This is what retrieval systems mean by
	//     "contextual chunking".
	//
	//     Note the prefix is a summary, not the whole head: prepending the
	//     entire head would just reconstruct the original sentence, which
	//     answers a different question (and defeats the point of chunking,
	//     since the budget is why we split in the first place).
	const contextPrefix = "About a movie."
	vPrefixed, err := embedUnit(contextPrefix + " " + tail)
	if err != nil {
		return err
	}
	fmt.Printf("     %.4f  (d) %q + tail — real attention\n", score(vPrefixed), contextPrefix)

	// For reference, the target we are trying to approximate.
	fmt.Printf("     %.4f  (--) full sentence, for comparison\n", score(vFull))

	fmt.Println()
	fmt.Println("   (b) and (c) are arithmetic on finished vectors: no token is")
	fmt.Println("   recomputed. (d) costs another forward pass, and is the only one")
	fmt.Println("   where the context reaches the tokens themselves.")
	fmt.Println("   Do not read too much into a single query: (b) scores well here")
	fmt.Println("   partly by dragging the vector toward the head, which happens to")
	fmt.Println("   mention films. The margin test below is the fairer comparison.")

	// The claim above is that (b) cannot fix a *wrong* chunk, only nudge it.
	// Test it directly on the violence query, where the tail is not merely
	// weak but misleading.
	// The discriminating question is not any single score but the *margin*:
	// does the representation rank the film query above the violence query,
	// the way the full sentence does? Absolute cosines are not comparable
	// across different queries, but the ordering within one representation is.
	vViolence, err := embedUnit(violenceQuery)
	if err != nil {
		return err
	}
	fmt.Println()
	fmt.Println("   The real test is the margin: does it prefer film over violence,")
	fmt.Println("   as the full sentence does?")
	fmt.Printf("     %-34s film %.4f  violence %.4f  margin %+.4f\n",
		"(a) tail alone", dot(vQuery, vTail), dot(vViolence, vTail),
		dot(vQuery, vTail)-dot(vViolence, vTail))
	fmt.Printf("     %-34s film %.4f  violence %.4f  margin %+.4f\n",
		"(b) mean(head, tail)", dot(vQuery, avg), dot(vViolence, avg),
		dot(vQuery, avg)-dot(vViolence, avg))
	fmt.Printf("     %-34s film %.4f  violence %.4f  margin %+.4f\n",
		"(d) context prefix + tail", dot(vQuery, vPrefixed), dot(vViolence, vPrefixed),
		dot(vQuery, vPrefixed)-dot(vViolence, vPrefixed))
	fmt.Printf("     %-34s film %.4f  violence %.4f  margin %+.4f\n",
		"(--) full sentence", dot(vQuery, vFull), dot(vViolence, vFull),
		dot(vQuery, vFull)-dot(vViolence, vFull))
	fmt.Println("   The tail alone has a negative margin: it ranks violence above")
	fmt.Println("   film, the opposite of the truth. Both repairs restore the sign.")

	// -----------------------------------------------------------------
	// 7. What pooling throws away, and why injection is tempting.
	// -----------------------------------------------------------------
	//
	// The model does compute a distinct vector per token, all of which are
	// already context-aware. Pooling collapses them to one and discards the
	// rest. Show that the per-token vectors genuinely differ by context: the
	// same word in two sentences gets two different vectors.
	section("7. Token vectors are already contextual — pooling discards that")
	simSame, err := tokenSimilarity(emb, "the movie was about a bank robbery", "the movie was about a bank holiday", "bank")
	if err != nil {
		return err
	}
	fmt.Printf("     %.4f  vector for %q in 'bank robbery' vs 'bank holiday'\n", simSame, "bank")
	fmt.Println("   Same token id, two different vectors, because attention already")
	fmt.Println("   mixed in the neighbours. That per-token conditioning is what a")
	fmt.Println("   split destroys. Averaging finished vectors cannot recreate it —")
	fmt.Println("   it can only move the result to roughly the right neighbourhood,")
	fmt.Println("   which for ranking is often enough, and sometimes is not.")

	return nil
}

// tokenSimilarity embeds two sentences and returns the cosine similarity
// between the contextual vectors of the first occurrence of a shared word.
//
// It relies on the two sentences tokenizing identically up to the position of
// interest, which is true for the pairs used here.
func tokenSimilarity(e *embedder, sentA, sentB, word string) (float32, error) {
	flatA, seqA, hidden, err := e.forward(sentA)
	if err != nil {
		return 0, err
	}
	flatB, seqB, _, err := e.forward(sentB)
	if err != nil {
		return 0, err
	}
	idA, err := tokenIndex(e, sentA, word, seqA)
	if err != nil {
		return 0, err
	}
	idB, err := tokenIndex(e, sentB, word, seqB)
	if err != nil {
		return 0, err
	}
	va := normalize(flatA[idA*hidden : (idA+1)*hidden])
	vb := normalize(flatB[idB*hidden : (idB+1)*hidden])
	return dot(va, vb), nil
}

// tokenIndex finds the position of word's token within the encoding of sent.
func tokenIndex(e *embedder, sent, word string, seqLen int) (int, error) {
	want := e.tokenizer.Encode(word)
	// Encode adds [CLS] and [SEP]; the bare token is in between.
	if len(want) < 3 {
		return 0, fmt.Errorf("word %q did not tokenize to a single piece", word)
	}
	target := want[1]
	got := e.tokenizer.Encode(sent)
	for i, id := range got {
		if id == target {
			return i, nil
		}
	}
	return 0, fmt.Errorf("token for %q not found in %q", word, sent)
}

func weightedMean(a, b []float32, wa, wb float32) []float32 {
	out := make([]float32, len(a))
	for i := range a {
		out[i] = a[i]*wa + b[i]*wb
	}
	return out
}

// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------

func section(title string) {
	fmt.Printf("\n%s\n%s\n", title, strings.Repeat("-", len(title)))
}

// compare prints the cosine similarity between two texts.
func compare(embed func(string) ([]float32, error), a, b string) error {
	va, err := embed(a)
	if err != nil {
		return err
	}
	vb, err := embed(b)
	if err != nil {
		return err
	}
	fmt.Printf("   %.4f  %q\n", dot(va, vb), a)
	fmt.Printf("           %q\n", b)
	return nil
}

// compareLabeled prints the similarity of text against a fixed query, with a
// short label, so several candidates can be lined up under one query.
func compareLabeled(embed func(string) ([]float32, error), query, text, label string) error {
	vq, err := embed(query)
	if err != nil {
		return err
	}
	vt, err := embed(text)
	if err != nil {
		return err
	}
	fmt.Printf("     %.4f  %s\n", dot(vq, vt), label)
	return nil
}

// normalize scales a vector to length 1, so that a dot product is the cosine
// of the angle between two of them. See DESIGN.md.
func normalize(v []float32) []float32 {
	var sum float64
	for _, x := range v {
		sum += float64(x) * float64(x)
	}
	n := float32(math.Sqrt(sum))
	if n == 0 {
		return v
	}
	out := make([]float32, len(v))
	for i, x := range v {
		out[i] = x / n
	}
	return out
}

func dot(a, b []float32) float32 {
	var sum float32
	for i := range a {
		sum += a[i] * b[i]
	}
	return sum
}

func mean(vs ...[]float32) []float32 {
	out := make([]float32, len(vs[0]))
	for _, v := range vs {
		for i, x := range v {
			out[i] += x
		}
	}
	for i := range out {
		out[i] /= float32(len(vs))
	}
	return out
}