kernel-code-embed / README.md
nethunter2023's picture
Add general-purpose baseline, paired significance test, and eval protocol
f3271f7 verified
|
Raw
History Blame Contribute Delete
4.42 kB
---
license: gpl-2.0
library_name: sentence-transformers
pipeline_tag: sentence-similarity
tags:
- sentence-transformers
- feature-extraction
- sentence-similarity
- code-retrieval
- code-search
- linux-kernel
- c
language:
- code
---
# kernel-code-embed
A 42.6M-parameter bi-encoder for retrieving **Linux kernel C** from
natural-language queries. Queries and code go through the same encoder, with no
prefix or instruction prompt.
512-dim output, mean pooling, L2-normalised β€” score with cosine similarity.
## Usage
```python
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("nethunter2023/kernel-code-embed")
query = "how are free pages coalesced into larger blocks"
code = "static inline void __free_one_page(struct page *page, ...) { ... }"
emb = model.encode([query, code], normalize_embeddings=True)
print(emb @ emb.T) # cosine similarity
```
### Sequence length
`max_seq_length` is **320**, and should not be raised even though
`max_position_embeddings` is 512. Pretraining saw 512 tokens but the retrieval
stage trained at 320, so positions 320–511 are undertrained; feeding longer
inputs measurably degrades ranking. Chunk longer functions instead.
## Results
Two different protocols. Read them separately β€” the candidate pools differ by
two orders of magnitude, so the numbers are **not** comparable across tables.
### 1. Open retrieval β€” the whole kernel
Every `.c`/`.h` chunk of Linux v7.1-rc5 as the index (**914,554 candidates**),
**N = 400** held-out queries, one correct answer each.
| | this model | + BM25 fusion |
|---|---|---|
| recall@1 | **0.8125** | 0.9050 |
| recall@5 | 0.9375 | 0.9725 |
| recall@10 | 0.9575 | 0.9775 |
| recall@50 | 0.9850 | 0.9950 |
| MRR | 0.8682 | 0.9374 |
**The model alone ranks the correct function first 81% of the time out of
914,554 candidates** (95% CI Β±3.8pp at N=400).
The second column is reciprocal-rank fusion of this model with BM25. It is
higher, but part of that lift is BM25's, so the 0.8125 figure is the one that
belongs to this model.
### 2. Closed set β€” against other encoders
**N = 2,000** held-out queries, 4,000 candidates, where each distractor is a
sibling function from the same source file as the answer.
| | params | accuracy@1 | NDCG@10 |
|---|---|---|---|
| BM25 (lexical) | β€” | 0.7115 | 0.8297 |
| `jina-embeddings-v2-base-code` | 161M | 0.9035 | 0.9541 |
| **this model** | **42.6M** | **0.9230** | **0.9661** |
Scored with the identical query set, candidate pool and metric code. The
general-purpose code embedder was given a 512-token budget β€” more than this
model's 320.
The +2.0pp accuracy@1 margin over `jina-embeddings-v2-base-code` is significant
under a paired McNemar test: p = 0.003, 95% CI on the paired difference
+0.7pp to +3.4pp (N = 2,000). A 42.6M domain model edging out a 161M
general-purpose one β€” at 3.8Γ— fewer parameters β€” is the result worth having.
## How the evaluation set was built
This matters for reading the numbers above.
Queries are **kernel-doc comments written by kernel developers**, not generated
by a language model β€” so they are not distribution-matched to the training data
by construction. They are held out from training. The **symbol name is stripped
from the query**: left in, `kmalloc_node` would appear in both the query and the
answer's signature and the task would collapse to identifier matching.
Training used InfoNCE over in-batch negatives (batch size 24, softmax scale
20.0), with hard negatives drawn as sibling functions from the same file.
## Limitations
- **Domain-specific.** Linux kernel C only. Not a general-purpose code or text
embedding model; do not expect transfer to other languages or codebases.
- **Short context** β€” 320 tokens. Long functions must be chunked.
- **kernel-doc phrasing.** Queries are developer-written documentation, which is
more precise than typical end-user questions. Expect lower accuracy on casual
or ambiguous phrasing.
- **Single held-out split**, no seed variance reported.
## License
GPL-2.0, matching the Linux kernel corpus it was trained on. Whether a GPL
training corpus propagates to model weights is legally unsettled; this is the
conservative reading, chosen deliberately rather than by default. If you need
different terms for commercial use, treat that as an open question to resolve
with counsel rather than an answered one.