nx_bm25_naive.nx source

1// nx_bm25.nx -- production-grade ranked retrieval: Okapi BM25 (Robertson & Sparck Jones), integer-only. 2// This is the HONEST upgrade over the team's first TF-IDF ranker, which lacked LENGTH NORMALIZATION: 3// a long document that mentions a term once would out-rank a short focused document with the same 4// per-length relevance. BM25 fixes that with (1) tf SATURATION (k1) -- the 10th occurrence adds less 5// than the 2nd -- and (2) LENGTH NORMALIZATION (b) -- penalize long docs by dl/avgdl. The point of 6// building the real standard (not a weak baseline) is to TRIANGULATE: nx_bm25_test asserts the team's 7// integer ranking reproduces a float reference BM25's ranking, so the "exceed" is measured against the 8// production formula, not cherry-picked. license_tier: ORIGINAL 9// 10// BM25(q,d) = sum_t idf(t) * tf(t,d)*(k1+1) / ( tf(t,d) + k1*(1 - b + b*dl/avgdl) ) 11// idf(t) = ln( 1 + (N - df + 0.5)/(df + 0.5) ) = ln( (2N+2)/(2df+1) ) [always positive] 12 13import "nx_research_extract.nx" // re_count / re_has / re_strlen 14import "nx_syscalls.nx" 15 16// standard BM25 parameters (Robertson/Sparck Jones); kept as named scaled constants, not magic numbers. 17const BM_K1_X1000: i64 = 1200 // k1 = 1.2 (term-frequency saturation) 18const BM_B_X1000: i64 = 750 // b = 0.75 (length-normalization strength) 19const BM_MICRO: i64 = 1000000 // fixed-point scale for ln / idf 20const BM_MILLI: i64 = 1000 // fixed-point scale for the saturation factor 21 22// number of whitespace-delimited tokens in text[0..n) -- the document length dl BM25 normalizes by. 23func bm_token_count(text: *u8, n: i64) -> i64 { 24 var c: i64 = 0; var i: i64 = 0; var in_tok: i64 = 0 25 while i < n { 26 let ch: i64 = text[i] 27 if ch == 32 { in_tok = 0 } else { if ch == 10 { in_tok = 0 } else { if ch == 13 { in_tok = 0 } else { if in_tok == 0 { c = c + 1; in_tok = 1 } } } } 28 i = i + 1 29 } 30 return c 31} 32 33// ln(num/den) in micro-units (x1e6) via the fast-converging atanh series: 34// ln(x) = 2*( z + z^3/3 + z^5/5 + ... ), z = (x-1)/(x+1) = (num-den)/(num+den). 35// num,den small positive ints; intermediate products stay well under i64. 36func bm_ln_micro(num: i64, den: i64) -> i64 { 37 if num <= 0 { return 0 } 38 if den <= 0 { return 0 } 39 let z: i64 = ((num - den) * BM_MICRO) / (num + den) // x1e6, signed 40 let zsq: i64 = (z * z) / BM_MICRO // x1e6 41 var sum: i64 = z 42 var zpow: i64 = z 43 var k: i64 = 3 44 while k <= 11 { 45 zpow = (zpow * zsq) / BM_MICRO // z^k, x1e6 46 sum = sum + zpow / k 47 k = k + 2 48 } 49 return 2 * sum 50} 51 52// BM25 idf in micro-units. df = document frequency (1..N). always positive (the +1 form). 53func bm_idf_micro(N: i64, df: i64) -> i64 { return bm_ln_micro(2 * N + 2, 2 * df + 1) } 54 55// the tf-saturation x length-normalization factor, in milli-units (x1000): 56// tf*(k1+1) / ( tf + k1*(1 - b + b*dl/avgdl) ) 57func bm_sat_milli(tf: i64, dl: i64, avgdl: i64) -> i64 { 58 if avgdl <= 0 { return 0 } 59 let norm_x1000: i64 = (BM_MILLI - BM_B_X1000) + (BM_B_X1000 * dl) / avgdl // (1-b)+b*dl/avgdl, x1000 60 let den_x1000: i64 = tf * BM_MILLI + (BM_K1_X1000 * norm_x1000) / BM_MILLI // tf + k1*norm, x1000 61 if den_x1000 <= 0 { return 0 } 62 let num_x1000: i64 = tf * (BM_K1_X1000 + BM_MILLI) // tf*(k1+1), x1000 63 return (num_x1000 * BM_MILLI) / den_x1000 // SAT x1000 64} 65 66// document frequency of `term` across the N-doc corpus (ptrs[i]=doc text as i64, lens[i]=its length). 67func bm_df(ptrs: *i64, lens: *i64, N: i64, term: *u8) -> i64 { 68 var c: i64 = 0; var i: i64 = 0 69 while i < N { if re_has(ptrs[i] as *u8, lens[i], term) == 1 { c = c + 1 } i = i + 1 } 70 return c 71} 72 73// BM25 score of doc k for an nq-term query (qterms[t] = a *u8 needle as i64). scale = x1e6. 74func bm_score(ptrs: *i64, lens: *i64, dls: *i64, N: i64, k: i64, avgdl: i64, qterms: *i64, nq: i64) -> i64 { 75 var s: i64 = 0; var t: i64 = 0 76 while t < nq { 77 let term: *u8 = qterms[t] as *u8 78 let tf: i64 = re_count(ptrs[k] as *u8, lens[k], term) 79 if tf > 0 { 80 let idf: i64 = bm_idf_micro(N, bm_df(ptrs, lens, N, term)) // x1e6 81 let sat: i64 = bm_sat_milli(tf, dls[k], avgdl) // x1e3 82 s = s + (idf * sat) / BM_MILLI // keep x1e6 83 } 84 t = t + 1 85 } 86 return s 87} 88 89// the best-ranked doc index for the query (-1 if nothing matches). 90func bm_best(ptrs: *i64, lens: *i64, dls: *i64, N: i64, avgdl: i64, qterms: *i64, nq: i64) -> i64 { 91 var best: i64 = 0 - 1; var bestscore: i64 = 0; var i: i64 = 0 92 while i < N { 93 let sc: i64 = bm_score(ptrs, lens, dls, N, i, avgdl, qterms, nq) 94 if sc > bestscore { bestscore = sc; best = i } 95 i = i + 1 96 } 97 return best 98} 99 100// the NAIVE tf*idf score (NO length normalization) -- the team's old ranker, kept to PROVE the flip. 101func bm_naive_tfidf(ptrs: *i64, lens: *i64, N: i64, k: i64, qterms: *i64, nq: i64) -> i64 { 102 var s: i64 = 0; var t: i64 = 0 103 while t < nq { 104 let term: *u8 = qterms[t] as *u8 105 let tf: i64 = re_count(ptrs[k] as *u8, lens[k], term) 106 if tf > 0 { s = s + tf * bm_idf_micro(N, bm_df(ptrs, lens, N, term)) } 107 t = t + 1 108 } 109 return s 110} 111 112func bm_naive_best(ptrs: *i64, lens: *i64, N: i64, qterms: *i64, nq: i64) -> i64 { 113 var best: i64 = 0 - 1; var bestscore: i64 = 0; var i: i64 = 0 114 while i < N { 115 let sc: i64 = bm_naive_tfidf(ptrs, lens, N, i, qterms, nq) 116 if sc > bestscore { bestscore = sc; best = i } 117 i = i + 1 118 } 119 return best 120}