Skip to content

Commit 916b9fa

Browse files
XiaoConstantineXiaoConstantine
committed
Make TQ-MSE the default ColBERT codec
1 parent 5b7fbbb commit 916b9fa

23 files changed

Lines changed: 2888 additions & 105 deletions

‎README.md‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,8 @@ sgrep setup
7878
# Index your codebase (vectors + ColBERT preindex, auto-starts embedding server)
7979
sgrep index .
8080

81-
# Optional: use the PQ6 ColBERT codec for larger repos
82-
# Small repos automatically stay on int8
81+
# Optional: override the default TQ-MSE ColBERT codec
82+
sgrep index . --colbert-codec int8
8383
sgrep index . --colbert-codec pq6
8484

8585
# Semantic search (quick)
@@ -513,6 +513,9 @@ Recent `dspy-go` benchmark on Apple M3 Pro + Metal (532 files, 7,735 chunks, 35,
513513
| Search quality | tuned pure `pq6` matched the current int8 benchmark at **MRR 0.725** |
514514
515515
Notes:
516+
- The default ColBERT segment codec is `tqmse`, which targets about half the
517+
segment storage of int8 while keeping int8 available as a conservative
518+
`--colbert-codec int8` override.
516519
- `--colbert-codec pq6` is size-gated: small repos automatically stay on int8.
517520
- Search latency varies significantly with corpus size, hardware, and whether ColBERT or rerank are enabled, so the README avoids claiming one universal query-time number.
518521

‎internal/cli/root.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -385,7 +385,7 @@ func init() {
385385
indexCmd.Flags().StringVar(&indexQuantize, "quantize", "int8", "Quantization mode: none (4x size), int8 (1x size), binary (0.125x size)")
386386
indexCmd.Flags().BoolVar(&indexColBERTPreindex, "colbert-preindex", true, "Pre-compute ColBERT segment embeddings for fast query-time scoring (default: true)")
387387
indexCmd.Flags().BoolVar(&indexColBERTAdaptiveSegments, "colbert-adaptive-segments", false, "Use adaptive sqrt(M) segment budgets during ColBERT preindexing (experimental)")
388-
indexCmd.Flags().StringVar(&indexColBERTCodec, "colbert-codec", "", "ColBERT segment codec: int8 or pq6 (default: reuse existing repo codec, otherwise int8)")
388+
indexCmd.Flags().StringVar(&indexColBERTCodec, "colbert-codec", "", "ColBERT segment codec: tqmse, int8, or pq6 (default: reuse existing repo codec, otherwise tqmse)")
389389
}
390390

391391
// Watch command

‎pkg/index/colbert_pq_scratch.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ func cloneColBERTSegments(segments []store.ColBERTSegment) []store.ColBERTSegmen
107107
Embedding: append([]float32(nil), seg.Embedding...),
108108
EmbeddingInt8: append([]int8(nil), seg.EmbeddingInt8...),
109109
PQCodes: append([]byte(nil), seg.PQCodes...),
110+
TQCodes: append([]byte(nil), seg.TQCodes...),
110111
}
111112
}
112113
return cloned

‎pkg/index/index.go‎

Lines changed: 148 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ type IndexConfig struct {
3434
Quantization store.QuantizationMode // Vector quantization mode (none, int8, binary)
3535
SmartSkip bool // Enable smart skipping for large repos (default: true)
3636
AdaptiveColBERTSegments bool // Enable token-aware sqrt(M) ColBERT segment budgets
37-
ColBERTCodec store.ColBERTCodec // Late interaction segment codec (int8, pq6)
37+
ColBERTCodec store.ColBERTCodec // Late interaction segment codec (tqmse, int8, pq6)
3838
}
3939

4040
// Large repo threshold - above this we enable smart skipping
@@ -47,10 +47,30 @@ const (
4747
colbertPQMinTrainingVectors = 256
4848
colbertPQTrainIterations = 25
4949
colbertPQMinWorthwhileSegments = 5000
50+
colbertTQMSEBits = 4
51+
colbertTQMSESeed = 42
5052
colbertPreindexChunkFetchBatchSize = 64
5153
colbertRewriteChunkBatchSize = 128
5254
)
5355

56+
func colbertEmbeddingDims() int {
57+
if v := os.Getenv("SGREP_DIMS"); v != "" {
58+
var dims int
59+
if _, err := fmt.Sscanf(v, "%d", &dims); err == nil && dims > 0 {
60+
return dims
61+
}
62+
}
63+
return 768
64+
}
65+
66+
func newDefaultColBERTTQMSEQuantizer() (*util.TQMSEQuantizer, error) {
67+
return util.NewTQMSEQuantizer(util.TQMSEConfig{
68+
Dims: colbertEmbeddingDims(),
69+
Bits: colbertTQMSEBits,
70+
Seed: colbertTQMSESeed,
71+
})
72+
}
73+
5474
// DefaultIndexConfig returns sensible defaults for indexing.
5575
func DefaultIndexConfig() *IndexConfig {
5676
// Workers = parallel file readers (CPU-bound: read + chunk)
@@ -86,6 +106,7 @@ type Indexer struct {
86106
ignore *IgnoreRules
87107
colbertCodec store.ColBERTCodec
88108
colbertPQ *util.ProductQuantizer
109+
colbertTQMSE *util.TQMSEQuantizer
89110
processed int64
90111
errors int64
91112
}
@@ -132,14 +153,19 @@ func NewWithConfig(path string, cfg *IndexConfig) (*Indexer, error) {
132153
}
133154
existingCodec := store.ColBERTCodecUnspecified
134155
var existingPQ *util.ProductQuantizer
156+
var existingTQMSE *util.TQMSEQuantizer
135157
if provider, ok := s.(store.ColBERTMetadataProvider); ok {
136158
existingCodec = provider.ColBERTCodec()
137159
existingPQ = provider.ProductQuantizer()
160+
existingTQMSE = provider.TQMSEQuantizer()
138161
}
139162
effectiveCodec := store.ResolveColBERTCodec(cfg.ColBERTCodec, existingCodec)
140163
if effectiveCodec != store.ColBERTCodecPQ6 {
141164
existingPQ = nil
142165
}
166+
if effectiveCodec != store.ColBERTCodecTQMSE {
167+
existingTQMSE = nil
168+
}
143169

144170
// Load ignore rules
145171
ignore := NewIgnoreRules(absPath)
@@ -154,6 +180,7 @@ func NewWithConfig(path string, cfg *IndexConfig) (*Indexer, error) {
154180
ignore: ignore,
155181
colbertCodec: effectiveCodec,
156182
colbertPQ: existingPQ,
183+
colbertTQMSE: existingTQMSE,
157184
}, nil
158185
}
159186

@@ -1553,6 +1580,7 @@ func (idx *Indexer) applyPQCodecSizeGate(ctx context.Context, segmentStore store
15531580
estimatedSegments, colbertPQMinWorthwhileSegments)
15541581
idx.colbertCodec = store.ColBERTCodecInt8
15551582
idx.colbertPQ = nil
1583+
idx.colbertTQMSE = nil
15561584
return nil
15571585
}
15581586

@@ -1727,6 +1755,7 @@ func (idx *Indexer) computeFreshPQSegmentsFromScratch(ctx context.Context, segme
17271755
fmt.Printf("Skipping PQ6: only %d segment vectors available, keeping int8\n", totalVectors)
17281756
idx.colbertCodec = store.ColBERTCodecInt8
17291757
idx.colbertPQ = nil
1758+
idx.colbertTQMSE = nil
17301759
written, err := idx.rewriteScratchColBERTSegments(ctx, segmentStore, scratchPath, totalChunks)
17311760
if err != nil {
17321761
return written, err
@@ -1754,6 +1783,7 @@ func (idx *Indexer) computeFreshPQSegmentsFromScratch(ctx context.Context, segme
17541783
fmt.Printf("ColBERT PQ train wall time: %v\n", time.Since(trainStart).Round(time.Millisecond))
17551784

17561785
idx.colbertPQ = pq
1786+
idx.colbertTQMSE = nil
17571787
written, err := idx.rewriteScratchColBERTSegments(ctx, segmentStore, scratchPath, totalChunks)
17581788
if err != nil {
17591789
return written, err
@@ -1787,6 +1817,9 @@ func (idx *Indexer) rewriteScratchColBERTSegments(ctx context.Context, segmentSt
17871817
if idx.colbertCodec == store.ColBERTCodecPQ6 && idx.colbertPQ != nil {
17881818
label = "Re-encoding ColBERT segments with PQ6"
17891819
progressVerb = "Re-encoded"
1820+
} else if idx.colbertCodec == store.ColBERTCodecTQMSE && idx.colbertTQMSE != nil {
1821+
label = "Encoding ColBERT segments with TQ-MSE"
1822+
progressVerb = "Encoded"
17901823
}
17911824
fmt.Printf("%s for %d chunks...\n", label, totalChunks)
17921825

@@ -1814,7 +1847,7 @@ func (idx *Indexer) rewriteScratchColBERTSegments(ctx context.Context, segmentSt
18141847
}
18151848
return processed, fmt.Errorf("read ColBERT PQ scratch record: %w", err)
18161849
}
1817-
converted, err := convertColBERTSegmentsForCodecStrict(record.ChunkID, record.Segments, idx.colbertCodec, idx.colbertPQ)
1850+
converted, err := convertColBERTSegmentsForCodecStrict(record.ChunkID, record.Segments, idx.colbertCodec, idx.colbertPQ, idx.colbertTQMSE)
18181851
if err != nil {
18191852
return processed, fmt.Errorf("convert scratch chunk %s to %s: %w", record.ChunkID, idx.colbertCodec, err)
18201853
}
@@ -1927,7 +1960,7 @@ func (idx *Indexer) ExportColBERTToMMap(ctx context.Context, outputDir string) (
19271960
defer func() { _ = os.RemoveAll(tempDir) }()
19281961

19291962
// Build the mmap in a temp location, then atomically replace the live file.
1930-
mmapStore, err := store.OpenMMapSegmentStore(tempDir, 768)
1963+
mmapStore, err := store.OpenMMapSegmentStore(tempDir, colbertEmbeddingDims())
19311964
if err != nil {
19321965
return 0, fmt.Errorf("failed to create MMap store: %w", err)
19331966
}
@@ -1936,8 +1969,21 @@ func (idx *Indexer) ExportColBERTToMMap(ctx context.Context, outputDir string) (
19361969
_ = mmapStore.Close()
19371970
}
19381971
}()
1972+
exportCodec := store.ColBERTCodecInt8
1973+
var exportPQ *util.ProductQuantizer
1974+
var exportTQ *util.TQMSEQuantizer
19391975
if provider, ok := segmentStore.(store.ColBERTMetadataProvider); ok {
1940-
mmapStore.SetColBERTCodec(provider.ColBERTCodec(), provider.ProductQuantizer())
1976+
switch provider.ColBERTCodec() {
1977+
case store.ColBERTCodecPQ6:
1978+
exportCodec = store.ColBERTCodecPQ6
1979+
exportPQ = provider.ProductQuantizer()
1980+
case store.ColBERTCodecTQMSE:
1981+
exportCodec = store.ColBERTCodecTQMSE
1982+
exportTQ = provider.TQMSEQuantizer()
1983+
default:
1984+
exportCodec = store.ColBERTCodecInt8
1985+
}
1986+
mmapStore.SetColBERTCodec(exportCodec, exportPQ, exportTQ)
19411987
}
19421988

19431989
mmapStore.BeginWrite()
@@ -1978,10 +2024,15 @@ func (idx *Indexer) ExportColBERTToMMap(ctx context.Context, outputDir string) (
19782024
return 0, fmt.Errorf("failed to load segments: %w", err)
19792025
}
19802026

1981-
// Write segments to MMap
2027+
// Write segments to MMap in the artifact codec. SQLite stores may be
2028+
// mixed after an upgrade or partial refresh.
19822029
for chunkID, chunkSegments := range segments {
1983-
mmapStore.WriteSegments(chunkID, chunkSegments)
1984-
totalSegments += len(chunkSegments)
2030+
encodedSegments, err := convertColBERTSegmentsForCodecStrict(chunkID, chunkSegments, exportCodec, exportPQ, exportTQ)
2031+
if err != nil {
2032+
return 0, fmt.Errorf("failed to encode chunk %s for MMap export: %w", chunkID, err)
2033+
}
2034+
mmapStore.WriteSegments(chunkID, encodedSegments)
2035+
totalSegments += len(encodedSegments)
19852036
}
19862037
}
19872038

@@ -2060,7 +2111,7 @@ func (idx *Indexer) buildColBERTChunkSegments(ctx context.Context, chunks []stor
20602111

20612112
chunkSegments := make(map[string][]store.ColBERTSegment, len(floatSegments))
20622113
for chunkID, segments := range floatSegments {
2063-
chunkSegments[chunkID] = convertColBERTSegmentsForCodec(chunkID, segments, idx.colbertCodec, idx.colbertPQ)
2114+
chunkSegments[chunkID] = convertColBERTSegmentsForCodec(chunkID, segments, idx.colbertCodec, idx.colbertPQ, idx.colbertTQMSE)
20642115
}
20652116
return chunkSegments, nil
20662117
}
@@ -2115,9 +2166,9 @@ func (idx *Indexer) buildFloat32ColBERTChunkSegments(ctx context.Context, chunks
21152166
return chunkSegments, nil
21162167
}
21172168

2118-
func buildStoredColBERTSegments(chunkID string, segmentTexts []string, embeddings [][]float32, adaptive bool, codec store.ColBERTCodec, pq *util.ProductQuantizer) []store.ColBERTSegment {
2169+
func buildStoredColBERTSegments(chunkID string, segmentTexts []string, embeddings [][]float32, adaptive bool, codec store.ColBERTCodec, pq *util.ProductQuantizer, tq *util.TQMSEQuantizer) []store.ColBERTSegment {
21192170
segments := buildFloat32ColBERTSegments(segmentTexts, embeddings, adaptive)
2120-
return convertColBERTSegmentsForCodec(chunkID, segments, codec, pq)
2171+
return convertColBERTSegmentsForCodec(chunkID, segments, codec, pq, tq)
21212172
}
21222173

21232174
func buildFloat32ColBERTSegments(segmentTexts []string, embeddings [][]float32, adaptive bool) []store.ColBERTSegment {
@@ -2147,27 +2198,38 @@ func buildFloat32ColBERTSegments(segmentTexts []string, embeddings [][]float32,
21472198
return segments
21482199
}
21492200

2150-
func convertColBERTSegmentsForCodec(chunkID string, segments []store.ColBERTSegment, codec store.ColBERTCodec, pq *util.ProductQuantizer) []store.ColBERTSegment {
2151-
converted, _ := convertColBERTSegmentsForCodecWithFallback(chunkID, segments, codec, pq, true)
2201+
func convertColBERTSegmentsForCodec(chunkID string, segments []store.ColBERTSegment, codec store.ColBERTCodec, pq *util.ProductQuantizer, tq *util.TQMSEQuantizer) []store.ColBERTSegment {
2202+
converted, _ := convertColBERTSegmentsForCodecWithFallback(chunkID, segments, codec, pq, tq, true)
21522203
return converted
21532204
}
21542205

2155-
func convertColBERTSegmentsForCodecStrict(chunkID string, segments []store.ColBERTSegment, codec store.ColBERTCodec, pq *util.ProductQuantizer) ([]store.ColBERTSegment, error) {
2156-
return convertColBERTSegmentsForCodecWithFallback(chunkID, segments, codec, pq, false)
2206+
func convertColBERTSegmentsForCodecStrict(chunkID string, segments []store.ColBERTSegment, codec store.ColBERTCodec, pq *util.ProductQuantizer, tq *util.TQMSEQuantizer) ([]store.ColBERTSegment, error) {
2207+
return convertColBERTSegmentsForCodecWithFallback(chunkID, segments, codec, pq, tq, false)
21572208
}
21582209

2159-
func convertColBERTSegmentsForCodecWithFallback(chunkID string, segments []store.ColBERTSegment, codec store.ColBERTCodec, pq *util.ProductQuantizer, allowPQFallback bool) ([]store.ColBERTSegment, error) {
2210+
func convertColBERTSegmentsForCodecWithFallback(chunkID string, segments []store.ColBERTSegment, codec store.ColBERTCodec, pq *util.ProductQuantizer, tq *util.TQMSEQuantizer, allowFallback bool) ([]store.ColBERTSegment, error) {
21602211
if codec == store.ColBERTCodecPQ6 && pq != nil {
21612212
encoded, err := encodeSegmentsToPQ(chunkID, segments, pq)
21622213
if err != nil {
2163-
if !allowPQFallback {
2214+
if !allowFallback {
21642215
return nil, err
21652216
}
21662217
fmt.Fprintf(os.Stderr, "Warning: ColBERT PQ encode failed for chunk %s, falling back to int8: %v\n", chunkID, err)
21672218
} else {
21682219
return encoded, nil
21692220
}
21702221
}
2222+
if codec == store.ColBERTCodecTQMSE && tq != nil {
2223+
encoded, err := encodeSegmentsToTQMSE(chunkID, segments, tq)
2224+
if err != nil {
2225+
if !allowFallback {
2226+
return nil, err
2227+
}
2228+
fmt.Fprintf(os.Stderr, "Warning: ColBERT TQ-MSE encode failed for chunk %s, falling back to int8: %v\n", chunkID, err)
2229+
} else {
2230+
return encoded, nil
2231+
}
2232+
}
21712233

21722234
for i := range segments {
21732235
if segments[i].Embedding != nil {
@@ -2181,15 +2243,67 @@ func convertColBERTSegmentsForCodecWithFallback(chunkID string, segments []store
21812243
return segments, nil
21822244
}
21832245

2246+
func encodeSegmentsToTQMSE(chunkID string, segments []store.ColBERTSegment, tq *util.TQMSEQuantizer) ([]store.ColBERTSegment, error) {
2247+
if tq == nil {
2248+
return nil, fmt.Errorf("nil tqmse quantizer")
2249+
}
2250+
encoded := make([]store.ColBERTSegment, len(segments))
2251+
encoder := tq.NewEncoder()
2252+
codeSize := tq.CodeSize()
2253+
for i, seg := range segments {
2254+
if len(seg.TQCodes) == codeSize {
2255+
encoded[i] = store.ColBERTSegment{
2256+
SegmentIdx: seg.SegmentIdx,
2257+
Text: seg.Text,
2258+
TQCodes: append([]byte(nil), seg.TQCodes...),
2259+
}
2260+
continue
2261+
}
2262+
emb := seg.Embedding
2263+
if emb == nil && len(seg.EmbeddingInt8) > 0 {
2264+
emb = util.DequantizeInt8(seg.EmbeddingInt8, seg.QuantScale, seg.QuantMin)
2265+
emb = util.NormalizeVector(emb)
2266+
}
2267+
if emb == nil {
2268+
if len(seg.TQCodes) > 0 {
2269+
return nil, fmt.Errorf("segment %d in chunk %s has malformed TQ-MSE code length %d (expected %d)", seg.SegmentIdx, chunkID, len(seg.TQCodes), codeSize)
2270+
}
2271+
return nil, fmt.Errorf("segment %d in chunk %s is missing an encodable embedding", seg.SegmentIdx, chunkID)
2272+
}
2273+
code, err := encoder.Encode(emb)
2274+
if err != nil {
2275+
return nil, err
2276+
}
2277+
encoded[i] = store.ColBERTSegment{
2278+
SegmentIdx: seg.SegmentIdx,
2279+
Text: seg.Text,
2280+
TQCodes: append([]byte(nil), code.Codes...),
2281+
}
2282+
}
2283+
return encoded, nil
2284+
}
2285+
21842286
func encodeSegmentsToPQ(chunkID string, segments []store.ColBERTSegment, pq *util.ProductQuantizer) ([]store.ColBERTSegment, error) {
21852287
encoded := make([]store.ColBERTSegment, len(segments))
2288+
codeSize := pq.CodeSize()
21862289
for i, seg := range segments {
2290+
if len(seg.PQCodes) == codeSize {
2291+
encoded[i] = store.ColBERTSegment{
2292+
SegmentIdx: seg.SegmentIdx,
2293+
Text: seg.Text,
2294+
PQCodes: append([]byte(nil), seg.PQCodes...),
2295+
}
2296+
continue
2297+
}
21872298
emb := seg.Embedding
21882299
if emb == nil && len(seg.EmbeddingInt8) > 0 {
21892300
emb = util.DequantizeInt8(seg.EmbeddingInt8, seg.QuantScale, seg.QuantMin)
21902301
emb = util.NormalizeVector(emb)
21912302
}
21922303
if emb == nil {
2304+
if len(seg.PQCodes) > 0 {
2305+
return nil, fmt.Errorf("segment %d in chunk %s has malformed PQ code length %d (expected %d)", seg.SegmentIdx, chunkID, len(seg.PQCodes), codeSize)
2306+
}
21932307
return nil, fmt.Errorf("segment %d in chunk %s is missing an encodable embedding", seg.SegmentIdx, chunkID)
21942308
}
21952309
codes, err := pq.Encode(emb)
@@ -2215,9 +2329,24 @@ func decomposeDocumentForColBERT(content string, adaptive bool) []string {
22152329
}
22162330

22172331
func (idx *Indexer) ensureColBERTCodecReady(ctx context.Context, segmentStore store.ColBERTSegmentStorer) error {
2332+
if idx.colbertCodec == store.ColBERTCodecTQMSE {
2333+
if idx.colbertTQMSE == nil {
2334+
tq, err := newDefaultColBERTTQMSEQuantizer()
2335+
if err != nil {
2336+
return fmt.Errorf("create colbert tqmse: %w", err)
2337+
}
2338+
idx.colbertTQMSE = tq
2339+
}
2340+
idx.colbertPQ = nil
2341+
return idx.persistColBERTMetadata(ctx)
2342+
}
2343+
22182344
if idx.colbertCodec != store.ColBERTCodecPQ6 {
2345+
idx.colbertPQ = nil
2346+
idx.colbertTQMSE = nil
22192347
return idx.persistColBERTMetadata(ctx)
22202348
}
2349+
idx.colbertTQMSE = nil
22212350
if idx.colbertPQ != nil {
22222351
return nil
22232352
}
@@ -2231,6 +2360,7 @@ func (idx *Indexer) ensureColBERTCodecReady(ctx context.Context, segmentStore st
22312360
fmt.Printf("Skipping PQ6: only %d segment vectors available, keeping int8\n", totalVectors)
22322361
idx.colbertCodec = store.ColBERTCodecInt8
22332362
idx.colbertPQ = nil
2363+
idx.colbertTQMSE = nil
22342364
return idx.persistColBERTMetadata(ctx)
22352365
}
22362366
fmt.Printf("Training ColBERT PQ codebook on %d sampled vectors (%d total)...\n", len(sample), totalVectors)
@@ -2249,6 +2379,7 @@ func (idx *Indexer) ensureColBERTCodecReady(ctx context.Context, segmentStore st
22492379
}
22502380

22512381
idx.colbertPQ = pq
2382+
idx.colbertTQMSE = nil
22522383
return nil
22532384
}
22542385

@@ -2257,7 +2388,7 @@ func (idx *Indexer) persistColBERTMetadata(ctx context.Context) error {
22572388
if !ok {
22582389
return nil
22592390
}
2260-
return metadataStore.SaveColBERTMetadata(ctx, idx.colbertCodec, idx.colbertPQ)
2391+
return metadataStore.SaveColBERTMetadata(ctx, idx.colbertCodec, idx.colbertPQ, idx.colbertTQMSE)
22612392
}
22622393

22632394
func (idx *Indexer) collectColBERTPQTrainingSample(ctx context.Context, segmentStore store.ColBERTSegmentStorer) ([][]float32, int, error) {

0 commit comments

Comments
 (0)