| 1 | // Retrieval V2, shadow-only: BM25F field weighting, code-symbol splitting, |
| 2 | // and mixed CJK uni+bigrams. Nothing here serves the model — V2 rankings ride |
| 3 | // the recall audit next to production's, and MemoryBench decides if V2 ever |
| 4 | // takes over. Weights are candidates under measurement, not tuned truths. |
| 5 | package retrieval |
| 6 | |
| 7 | import ( |
| 8 | "math" |
| 9 | "strings" |
| 10 | "unicode" |
| 11 | ) |
| 12 | |
| 13 | // TokensV2 extends Tokens with code-symbol structure and CJK unigrams: |
| 14 | // CamelCase and snake_case words also emit their segments (FooBar → foobar, |
| 15 | // foo, bar), and CJK runs emit unigrams alongside bigrams so one-character |
| 16 | // queries still match. The V1 form of every token is preserved, so V2 recall |
| 17 | // is a superset of V1's vocabulary. |
| 18 | func TokensV2(s string) []string { |
| 19 | var out []string |
| 20 | var word []rune |
| 21 | var prev rune |
| 22 | cjkRun := 0 |
| 23 | flushWord := func() { |
| 24 | if len(word) == 0 { |
| 25 | return |
| 26 | } |
| 27 | lower := strings.ToLower(string(word)) |
| 28 | out = append(out, lower) |
| 29 | for _, seg := range splitCodeSymbol(word) { |
| 30 | if seg != lower { |
| 31 | out = append(out, seg) |
| 32 | } |
| 33 | } |
| 34 | word = word[:0] |
| 35 | } |
| 36 | endCJK := func() { cjkRun = 0 } |
| 37 | for _, r := range s { |
| 38 | switch { |
| 39 | case isCJK(r): |
| 40 | flushWord() |
| 41 | out = append(out, string(r)) |
| 42 | if cjkRun > 0 { |
| 43 | out = append(out, string([]rune{prev, r})) |
| 44 | } |
| 45 | prev = r |
| 46 | cjkRun++ |
| 47 | case unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_': |
| 48 | endCJK() |
| 49 | word = append(word, r) |
| 50 | default: |
| 51 | flushWord() |
| 52 | endCJK() |
| 53 | } |
| 54 | } |
| 55 | flushWord() |
| 56 | return out |
| 57 | } |
| 58 | |
| 59 | // splitCodeSymbol yields the lowercase segments of a code identifier: |
| 60 | // CamelCase humps, snake_case parts, and letter/digit boundaries. |
| 61 | func splitCodeSymbol(word []rune) []string { |
| 62 | var segments []string |
| 63 | var seg []rune |
| 64 | flush := func() { |
| 65 | if len(seg) > 1 { // single letters are noise, not signal |
| 66 | segments = append(segments, strings.ToLower(string(seg))) |
| 67 | } |
| 68 | seg = seg[:0] |
| 69 | } |
| 70 | for i, r := range word { |
| 71 | boundary := r == '_' || |
| 72 | (i > 0 && unicode.IsUpper(r) && !unicode.IsUpper(word[i-1])) || |
| 73 | (i > 0 && unicode.IsDigit(r) != unicode.IsDigit(word[i-1])) |
| 74 | if boundary { |
| 75 | flush() |
| 76 | } |
| 77 | if r != '_' { |
| 78 | seg = append(seg, r) |
| 79 | } |
| 80 | } |
| 81 | flush() |
| 82 | if len(segments) < 2 { |
| 83 | return nil // no structure worth indexing beyond the whole word |
| 84 | } |
| 85 | return segments |
| 86 | } |
| 87 | |
| 88 | // FieldedDoc is one document split into weighted fields for BM25F. |
| 89 | type FieldedDoc struct { |
| 90 | ID string |
| 91 | Fields map[string]string |
| 92 | } |
| 93 | |
| 94 | // v2FieldWeights follow the BM25F intuition: identity fields say what a fact |
| 95 | // IS, the body says everything it mentions. Candidates under measurement. |
| 96 | var v2FieldWeights = map[string]float64{ |
| 97 | "name": 2.5, "title": 2.5, "keywords": 2.0, "subject": 2.0, |
| 98 | "description": 1.5, "body": 1.0, |
| 99 | } |
| 100 | |
| 101 | // V2Hit is one shadow-ranked document. |
| 102 | type V2Hit struct { |
| 103 | ID string |
| 104 | Score float64 |
| 105 | } |
| 106 | |
| 107 | // RankV2 scores docs against a query with BM25F over TokensV2: term |
| 108 | // frequencies accumulate per field scaled by field weight, then the standard |
| 109 | // saturation applies once per term (the F in BM25F), with length |
| 110 | // normalization over the weighted document length. |
| 111 | func RankV2(query string, docs []FieldedDoc) []V2Hit { |
| 112 | queryTerms := Unique(TokensV2(query)) |
| 113 | if len(queryTerms) == 0 || len(docs) == 0 { |
| 114 | return nil |
| 115 | } |
| 116 | type indexed struct { |
| 117 | id string |
| 118 | tf map[string]float64 |
| 119 | length float64 |
| 120 | } |
| 121 | corpus := make([]indexed, 0, len(docs)) |
| 122 | df := map[string]int{} |
| 123 | var totalLen float64 |
| 124 | for _, doc := range docs { |
| 125 | ix := indexed{id: doc.ID, tf: map[string]float64{}} |
| 126 | seen := map[string]bool{} |
| 127 | for field, text := range doc.Fields { |
| 128 | weight, ok := v2FieldWeights[field] |
| 129 | if !ok { |
| 130 | weight = 1.0 |
| 131 | } |
| 132 | for _, term := range TokensV2(text) { |
| 133 | ix.tf[term] += weight |
| 134 | ix.length += weight |
| 135 | seen[term] = true |
| 136 | } |
| 137 | } |
| 138 | for term := range seen { |
| 139 | df[term]++ |
| 140 | } |
| 141 | totalLen += ix.length |
| 142 | corpus = append(corpus, ix) |
| 143 | } |
| 144 | avgLen := totalLen / float64(len(corpus)) |
| 145 | if avgLen <= 0 { |
| 146 | avgLen = 1 |
| 147 | } |
| 148 | const k1, b = 1.2, 0.75 |
| 149 | var hits []V2Hit |
| 150 | for _, doc := range corpus { |
| 151 | var score float64 |
| 152 | for _, term := range queryTerms { |
| 153 | tf := doc.tf[term] |
| 154 | if tf == 0 || df[term] == 0 { |
| 155 | continue |
| 156 | } |
| 157 | idf := math.Log(1 + (float64(len(corpus))-float64(df[term])+0.5)/(float64(df[term])+0.5)) |
| 158 | score += idf * (tf * (k1 + 1)) / (tf + k1*(1-b+b*doc.length/avgLen)) |
| 159 | } |
| 160 | if score > 0 { |
| 161 | hits = append(hits, V2Hit{ID: doc.id, Score: score}) |
| 162 | } |
| 163 | } |
| 164 | SortHitsDesc(hits) |
| 165 | return hits |
| 166 | } |
| 167 | |
| 168 | // SortHitsDesc orders shadow hits best-first with a stable ID tiebreak. |
| 169 | func SortHitsDesc(hits []V2Hit) { |
| 170 | for i := 1; i < len(hits); i++ { |
| 171 | for j := i; j > 0 && (hits[j].Score > hits[j-1].Score || |
| 172 | (hits[j].Score == hits[j-1].Score && hits[j].ID < hits[j-1].ID)); j-- { |
| 173 | hits[j], hits[j-1] = hits[j-1], hits[j] |
| 174 | } |
| 175 | } |
| 176 | } |
| 177 |