kernel-code-embed

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+  "word_embedding_dimension": 512,
+  "pooling_mode_cls_token": false,
+  "pooling_mode_mean_tokens": true,
+  "pooling_mode_max_tokens": false,
+  "pooling_mode_mean_sqrt_len_tokens": false,
+  "pooling_mode_weightedmean_tokens": false,
+  "pooling_mode_lasttoken": false,
+  "include_prompt": true
+}

README.md

@@ -0,0 +1,107 @@
+---
+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 compact **C code embedding model for Linux kernel retrieval**. Ask a question
+in English; it ranks kernel functions by how well they answer it.
+
+Symmetric bi-encoder — queries and code go through the **same** encoder with no
+prefix or instruction prompt.
+
+| | |
+|---|---|
+| parameters | 42.6M |
+| layers / hidden / heads | 8 / 512 / 8 |
+| vocabulary | 32,768 byte-level BPE, built for C |
+| max sequence length | **320** |
+| pooling | mean |
+| output | 512-dim, L2-normalised |
+| size on disk | 162 MB |
+
+## 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, unsigned long pfn,
+                                   struct zone *zone, unsigned int order,
+                                   int migratetype, fpi_t fpi_flags)
+{
+        ...
+}
+"""
+
+emb = model.encode([query, code], normalize_embeddings=True)
+print(emb @ emb.T)          # cosine similarity
+```
+
+**Do not raise `max_seq_length` above 320.** Feeding longer inputs measurably
+degrades retrieval quality — in a spot check it moved a correct answer from
+rank 12 to rank 40. Chunk longer functions instead.
+
+## Evaluating this model
+
+Enough detail to reproduce or compare the numbers below.
+
+- **Embeddings:** `normalize_embeddings=True`, mean pooling, 512-dim.
+- **Similarity:** cosine (equivalently, dot product on normalised vectors).
+- **No prompt prefixes** on either side — encode raw query text and raw code.
+- **Sequence length:** leave at 320.
+- **Open retrieval below:** every `.c`/`.h` chunk in Linux v7.1-rc5 as the index
+  (914,554 chunks), 400 held-out kernel-doc anchors as queries, the documented
+  function as the single correct answer.
+- **Hybrid** is dense retrieval fused with BM25 via reciprocal rank fusion.
+- **Closed evaluation:** 2,000 held-out pairs against 4,000 candidates. A
+  deliberately hard setting: the distractors come from the same source files as
+  the answer, so lexical overlap alone does not separate them.
+
+## Results
+
+Open retrieval over the whole kernel — **914,554 candidate chunks**:
+
+| metric | dense | hybrid (dense + BM25 RRF) |
+|---|---|---|
+| 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** |
+| median rank | 1 | 1 |
+
+The correct function ranks first out of 914,554 candidates 90% of the time.
+
+Closed evaluation, against a lexical baseline:
+
+| metric | BM25 | this model |
+|---|---|---|
+| accuracy@1 | 0.7115 | **0.9225** |
+| NDCG@10 | 0.8297 | **0.9659** |
+
+## Limitations
+
+- **Domain-specific.** Trained on Linux kernel C only. It is not a general-purpose
+  code or text embedding model, and it should not be expected to transfer to
+  other languages or codebases.
+- **Short context.** 320 tokens. Long functions must be chunked.
+- **Evaluated on kernel-doc anchors**, which are written by kernel developers and
+  are more precise than typical end-user questions. Expect lower accuracy on
+  casual or ambiguous phrasing.
+- Licensed GPL-2.0, consistent with its Linux kernel training data.

config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "artifacts/bert-kernel-mlm/phase2",
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 8,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.47.1",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 32768
+}

config_sentence_transformers.json

@@ -0,0 +1,10 @@
+{
+  "__version__": {
+    "sentence_transformers": "3.3.1",
+    "transformers": "4.47.1",
+    "pytorch": "2.5.1+cu121"
+  },
+  "prompts": {},
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cef78bc1945c263922a5d7d451b1f72d7c4c38009d89751fa122cdcf75ccf336
+size 170105544

modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 320,
+  "do_lower_case": false
+}

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,53 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 320,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}