feat: validate encoder identity and expose retrieval modes
This commit is contained in:
@@ -2,7 +2,7 @@
|
||||
|
||||
A small local code retrieval model and a tool for giving a larger coding model useful source context.
|
||||
|
||||
Micro-scout indexes a repository, combines lexical and neural search, and returns **verified file paths, line ranges, and bounded source snippets**. Its MCP server keeps the model in memory between requests.
|
||||
Micro-scout indexes a repository, supports neural, lexical, and hybrid search, and returns **verified file paths, line ranges, and bounded source snippets**. Its MCP server keeps the model in memory between requests.
|
||||
|
||||
## Version 0.1
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@ Version 0.1 fine-tunes `sentence-transformers/all-MiniLM-L6-v2` as a **shared bi
|
||||
|
||||
For each batch, the query's paired function is the labeled positive. Other functions in the batch are treated as negatives. The loss averages query-to-code and code-to-query cross-entropy. A temperature scales cosine scores. The model learns ranking without generating text.
|
||||
|
||||
At inference, the repository's vectors are already in the index. Each request encodes only its query, compares vectors, and combines rankings with BM25. This makes persistent inference practical without Ollama or vLLM.
|
||||
At inference, the repository's vectors are already in the index. Each request encodes only its query and compares vectors. Optional hybrid mode combines rankings with BM25. This makes persistent inference practical without Ollama or vLLM. Dense mode is the default because it performed best on complete validation; the fusion weights were not tuned.
|
||||
|
||||
## Reproduce the local run
|
||||
|
||||
@@ -73,7 +73,7 @@ OPENBLAS_NUM_THREADS=1 uv run --no-sync micro-scout benchmark \
|
||||
--output runs/minilm-v1/latency.json
|
||||
```
|
||||
|
||||
This measures the warm harness, including search, context assembly, and checking returned files. Startup is reported separately. Five repeated development queries are used; these are not representative production traffic. Index construction is also reported separately.
|
||||
This measures the warm harness, including search, context assembly, and checking returned files. `startup_ms` covers model/index initialization after the core imports, not the entire interpreter startup. Five repeated development queries are used; these are not representative production traffic. Index construction is also reported separately.
|
||||
|
||||
## What these results cannot establish
|
||||
|
||||
|
||||
@@ -4,6 +4,8 @@
|
||||
|
||||
`scout_search` accepts a natural-language query, `top_k` from 1 to 50, and `max_chars` from 100 to 100,000. It returns ranked snippets and up to three graph neighbors within the shared source-character budget. Overlapping source lines are not repeated. Scores are not probabilities.
|
||||
|
||||
With a model, the default `mode` is `dense`, selected from the complete validation results. `lexical` uses BM25, which is useful for exact identifiers, and `hybrid` combines the two rankings. Without a model the CLI and MCP server default to `lexical`. The CLI exposes the same choice as `--mode`.
|
||||
|
||||
Markdown is excluded from search by default; set `include_docs=true` to include it. An optional `language` filter narrows candidates before rank fusion, for example to `python` or `typescript`. CLI equivalents are `--include-docs` and `--language python`.
|
||||
|
||||
Each snippet includes an opaque symbol ID, repository-relative path, line range, source text, file SHA-256, and verification/truncation flags. The response also identifies the index snapshot and model fingerprint. Truncation occurs at line boundaries.
|
||||
@@ -40,6 +42,8 @@ The server uses the official [MCP Python SDK v1 maintenance line](https://py.sdk
|
||||
|
||||
Indexing honors Git's ignored-file rules when available and skips hidden directories, common dependency/build directories, symlinks, unsupported extensions, and files above 1 MB. Python symbols use AST boundaries. Other supported text/code formats use 60-line chunks. Incomplete Python falls back to chunks.
|
||||
|
||||
Python functions longer than 48 lines also get 32-line search fragments with an 8-line overlap. The full symbol remains available through each fragment's `parent_id`. This gives the retriever access to code beyond the beginning of a long function. Neural input is still capped at 256 tokens per candidate; very long individual lines can exceed that representation budget. BM25 searches the complete candidate text.
|
||||
|
||||
An index is one SQLite file, atomically replaced after a successful build. A model fingerprint prevents combining incompatible query and code embeddings. Reindexing with unchanged weights reuses embeddings for unchanged normalized code.
|
||||
|
||||
Returned files are checked against their indexed content hashes. Modified or missing files are excluded with warnings; reads of stale IDs fail. New files require reindexing. This is a snapshot workflow, not a filesystem watcher. Restart the server after rebuilding an index so it loads the new snapshot.
|
||||
|
||||
@@ -32,6 +32,7 @@ def main() -> None:
|
||||
if name == "benchmark":
|
||||
command.add_argument("--iterations", type=int, default=100)
|
||||
command.add_argument("--output", type=Path)
|
||||
command.add_argument("--mode", choices=["lexical", "dense", "hybrid"], default=None)
|
||||
for command in (index_parser, *(sub.choices[n] for n in ("search", "serve", "benchmark"))):
|
||||
command.add_argument("--model", help="Local weights directory or Hugging Face model ID")
|
||||
command.add_argument("--device", default="cpu")
|
||||
@@ -62,7 +63,7 @@ def main() -> None:
|
||||
args.query,
|
||||
top_k=args.top_k,
|
||||
max_chars=args.max_chars,
|
||||
mode=args.mode or ("hybrid" if encoder else "lexical"),
|
||||
mode=args.mode or ("dense" if encoder else "lexical"),
|
||||
expand=not args.no_expand,
|
||||
language=args.language,
|
||||
include_docs=args.include_docs,
|
||||
@@ -80,7 +81,7 @@ def main() -> None:
|
||||
"remove documentation strings from Python functions",
|
||||
"combine lexical and neural search rankings",
|
||||
]
|
||||
mode = "hybrid" if encoder else "lexical"
|
||||
mode = args.mode or ("dense" if encoder else "lexical")
|
||||
for query in queries:
|
||||
scout.search(query, mode=mode)
|
||||
times = []
|
||||
|
||||
@@ -61,6 +61,17 @@ class Encoder:
|
||||
self.dimension = self.model.config.hidden_size
|
||||
self.parameter_count = sum(p.numel() for p in self.model.parameters())
|
||||
signature = hashlib.sha256(json.dumps(self.config, sort_keys=True).encode())
|
||||
# Weights alone do not identify an encoder: tokenization and model
|
||||
# configuration can change the vectors without changing the weights.
|
||||
tokenizer_config = json.loads(self.tokenizer.backend_tokenizer.to_str())
|
||||
for runtime_option in ("padding", "truncation"):
|
||||
tokenizer_config.pop(runtime_option, None)
|
||||
model_config = self.model.config.to_dict()
|
||||
for provenance in ("_name_or_path", "transformers_version"):
|
||||
model_config.pop(provenance, None)
|
||||
signature.update(json.dumps(tokenizer_config, sort_keys=True).encode())
|
||||
signature.update(json.dumps(self.tokenizer.special_tokens_map, sort_keys=True).encode())
|
||||
signature.update(json.dumps(model_config, sort_keys=True).encode())
|
||||
if local:
|
||||
weights = sorted(Path(model).glob("*.safetensors"))
|
||||
if not weights:
|
||||
|
||||
@@ -65,6 +65,8 @@ def evaluate(
|
||||
"candidates_per_query": len(rows),
|
||||
"repositories": len({r["repo"] for r in rows}),
|
||||
"dataset_sha256": dataset_hash,
|
||||
"evaluator_source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
|
||||
"fusion": {"dense_weight": 0.5, "per_method_limit": 100, "k": 60},
|
||||
"evaluation_ids_sha256": hashlib.sha256(
|
||||
json.dumps([r["id"] for r in rows]).encode()
|
||||
).hexdigest(),
|
||||
@@ -115,6 +117,9 @@ def evaluate(
|
||||
report["trained_hybrid_vs_bm25_mrr"] = paired_mrr_interval(
|
||||
ranks["trained_hybrid"], ranks["bm25"], [r["repo"] for r in rows]
|
||||
)
|
||||
report["trained_dense_vs_bm25_mrr"] = paired_mrr_interval(
|
||||
ranks["trained_dense"], ranks["bm25"], [r["repo"] for r in rows]
|
||||
)
|
||||
report["elapsed_seconds"] = time.perf_counter() - started
|
||||
atomic_json(output / f"{split}-metrics.json", report)
|
||||
write_jsonl(
|
||||
|
||||
@@ -88,6 +88,7 @@ class Scout:
|
||||
"path": symbol.path,
|
||||
"name": symbol.name,
|
||||
"kind": symbol.kind,
|
||||
"parent_id": symbol.parent,
|
||||
"language": symbol.language,
|
||||
"start_line": symbol.start_line,
|
||||
"end_line": end,
|
||||
@@ -104,7 +105,7 @@ class Scout:
|
||||
*,
|
||||
top_k: int = 6,
|
||||
max_chars: int = 12_000,
|
||||
mode: str = "hybrid",
|
||||
mode: str = "dense",
|
||||
expand: bool = True,
|
||||
language: str | None = None,
|
||||
include_docs: bool = False,
|
||||
|
||||
@@ -30,18 +30,21 @@ def create_server(scout: Scout):
|
||||
max_chars: int = 12_000,
|
||||
language: str | None = None,
|
||||
include_docs: bool = False,
|
||||
mode: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Find relevant code. English queries are the evaluated language. Returns verified paths,
|
||||
line ranges, bounded source text, and approximate graph neighbors. max_chars counts source
|
||||
characters, not model tokens. Markdown is excluded unless include_docs is true.
|
||||
Optional language filter examples: python, typescript, rust.
|
||||
mode: dense (default with a model), lexical, or hybrid. Dense performed best
|
||||
on held-out English descriptions. Try lexical for an exact identifier.
|
||||
No code is executed. No source files are changed.
|
||||
"""
|
||||
return scout.search(
|
||||
query,
|
||||
top_k=top_k,
|
||||
max_chars=max_chars,
|
||||
mode="hybrid" if scout.encoder else "lexical",
|
||||
mode=mode or ("dense" if scout.encoder else "lexical"),
|
||||
language=language,
|
||||
include_docs=include_docs,
|
||||
)
|
||||
|
||||
@@ -167,6 +167,16 @@ def parse_source(relative: str, text: str, *, chunk_lines: int = 60) -> list[Sym
|
||||
walk(child, prefix, parent)
|
||||
|
||||
walk(tree)
|
||||
# Long functions need searchable windows beyond the encoder's first
|
||||
# 256 tokens. Keep the full symbol so callers can read its context.
|
||||
for symbol in list(symbols):
|
||||
if symbol.kind != "function" or symbol.end_line - symbol.start_line + 1 <= 48:
|
||||
continue
|
||||
for start in range(symbol.start_line, symbol.end_line + 1, 24):
|
||||
end = min(start + 31, symbol.end_line)
|
||||
add(f"{symbol.name}::<{start}-{end}>", "fragment", start, end, symbol.id)
|
||||
if end == symbol.end_line:
|
||||
break
|
||||
covered = set()
|
||||
for symbol in symbols:
|
||||
covered.update(range(symbol.start_line, symbol.end_line + 1))
|
||||
|
||||
@@ -119,3 +119,18 @@ def test_training_updates_weights_and_can_resume(tiny_model, tmp_path):
|
||||
after.save(output / "last")
|
||||
with pytest.raises(ValueError, match="weights and optimizer state"):
|
||||
train(data, output, config, "cpu", output / "last")
|
||||
|
||||
|
||||
def test_fingerprint_detects_tokenizer_change(tiny_model):
|
||||
original = Encoder(str(tiny_model), max_length=32, query_length=16)
|
||||
path = tiny_model / "tokenizer.json"
|
||||
config = json.loads(path.read_text())
|
||||
vocab = config["model"]["vocab"]
|
||||
vocab["read"], vocab["write"] = vocab["write"], vocab["read"]
|
||||
path.write_text(json.dumps(config))
|
||||
changed = Encoder(str(tiny_model), max_length=32, query_length=16)
|
||||
assert changed.fingerprint != original.fingerprint
|
||||
assert (
|
||||
changed.tokenize(["read file"])["input_ids"].tolist()
|
||||
!= original.tokenize(["read file"])["input_ids"].tolist()
|
||||
)
|
||||
|
||||
@@ -263,3 +263,20 @@ def test_source_that_grows_after_indexing_is_bounded(repository, tmp_path):
|
||||
(repository / "numbers.py").write_text("x" * 1_000_001)
|
||||
with pytest.raises(ValueError, match="file-size limit"):
|
||||
scout.read(symbol.id)
|
||||
|
||||
|
||||
def test_long_function_tail_has_a_searchable_fragment(repository, tmp_path):
|
||||
lines = ["def long_function():"] + [f" value_{i} = {i}" for i in range(90)]
|
||||
lines.append(" return unique_tail_marker")
|
||||
(repository / "long.py").write_text("\n".join(lines) + "\n")
|
||||
path = tmp_path / "index.sqlite"
|
||||
build_index(repository, path)
|
||||
scout = Scout(Index(path))
|
||||
result = scout.search("unique_tail_marker", mode="lexical", top_k=1)
|
||||
hit = result["results"][0]
|
||||
assert hit["kind"] == "fragment"
|
||||
assert "unique_tail_marker" in hit["content"]
|
||||
assert hit["start_line"] > 48
|
||||
parent = scout.read(hit["parent_id"])
|
||||
assert parent["name"] == "long_function"
|
||||
assert parent["start_line"] == 1
|
||||
|
||||
Reference in New Issue
Block a user