JevCodeLocator-0.6B

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1.  Definitions.
+
+    "License" shall mean the terms and conditions for use, reproduction,
+    and distribution as defined by Sections 1 through 9 of this document.
+
+    "Licensor" shall mean the copyright owner or entity authorized by
+    the copyright owner that is granting the License.
+
+    "Legal Entity" shall mean the union of the acting entity and all
+    other entities that control, are controlled by, or are under common
+    control with that entity. For the purposes of this definition,
+    "control" means (i) the power, direct or indirect, to cause the
+    direction or management of such entity, whether by contract or
+    otherwise, or (ii) ownership of fifty percent (50%) or more of the
+    outstanding shares, or (iii) beneficial ownership of such entity.
+
+    "You" (or "Your") shall mean an individual or Legal Entity
+    exercising permissions granted by this License.
+
+    "Source" form shall mean the preferred form for making modifications,
+    including but not limited to software source code, documentation
+    source, and configuration files.
+
+    "Object" form shall mean any form resulting from mechanical
+    transformation or translation of a Source form, including but
+    not limited to compiled object code, generated documentation,
+    and conversions to other media types.
+
+    "Work" shall mean the work of authorship, whether in Source or
+    Object form, made available under the License, as indicated by a
+    copyright notice that is included in or attached to the work
+    (an example is provided in the Appendix below).
+
+    "Derivative Works" shall mean any work, whether in Source or Object
+    form, that is based on (or derived from) the Work and for which the
+    editorial revisions, annotations, elaborations, or other modifications
+    represent, as a whole, an original work of authorship. For the purposes
+    of this License, Derivative Works shall not include works that remain
+    separable from, or merely link (or bind by name) to the interfaces of,
+    the Work and Derivative Works thereof.
+
+    "Contribution" shall mean any work of authorship, including
+    the original version of the Work and any modifications or additions
+    to that Work or Derivative Works thereof, that is intentionally
+    submitted to Licensor for inclusion in the Work by the copyright owner
+    or by an individual or Legal Entity authorized to submit on behalf of
+    the copyright owner. For the purposes of this definition, "submitted"
+    means any form of electronic, verbal, or written communication sent
+    to the Licensor or its representatives, including but not limited to
+    communication on electronic mailing lists, source code control systems,
+    and issue tracking systems that are managed by, or on behalf of, the
+    Licensor for the purpose of discussing and improving the Work, but
+    excluding communication that is conspicuously marked or otherwise
+    designated in writing by the copyright owner as "Not a Contribution."
+
+    "Contributor" shall mean Licensor and any individual or Legal Entity
+    on behalf of whom a Contribution has been received by Licensor and
+    subsequently incorporated within the Work.
+
+2.  Grant of Copyright License. Subject to the terms and conditions of
+    this License, each Contributor hereby grants to You a perpetual,
+    worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+    copyright license to reproduce, prepare Derivative Works of,
+    publicly display, publicly perform, sublicense, and distribute the
+    Work and such Derivative Works in Source or Object form.
+
+3.  Grant of Patent License. Subject to the terms and conditions of
+    this License, each Contributor hereby grants to You a perpetual,
+    worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+    (except as stated in this section) patent license to make, have made,
+    use, offer to sell, sell, import, and otherwise transfer the Work,
+    where such license applies only to those patent claims licensable
+    by such Contributor that are necessarily infringed by their
+    Contribution(s) alone or by combination of their Contribution(s)
+    with the Work to which such Contribution(s) was submitted. If You
+    institute patent litigation against any entity (including a
+    cross-claim or counterclaim in a lawsuit) alleging that the Work
+    or a Contribution incorporated within the Work constitutes direct
+    or contributory patent infringement, then any patent licenses
+    granted to You under this License for that Work shall terminate
+    as of the date such litigation is filed.
+
+4.  Redistribution. You may reproduce and distribute copies of the
+    Work or Derivative Works thereof in any medium, with or without
+    modifications, and in Source or Object form, provided that You
+    meet the following conditions:
+
+    (a) You must give any other recipients of the Work or
+    Derivative Works a copy of this License; and
+
+    (b) You must cause any modified files to carry prominent notices
+    stating that You changed the files; and
+
+    (c) You must retain, in the Source form of any Derivative Works
+    that You distribute, all copyright, patent, trademark, and
+    attribution notices from the Source form of the Work,
+    excluding those notices that do not pertain to any part of
+    the Derivative Works; and
+
+    (d) If the Work includes a "NOTICE" text file as part of its
+    distribution, then any Derivative Works that You distribute must
+    include a readable copy of the attribution notices contained
+    within such NOTICE file, excluding those notices that do not
+    pertain to any part of the Derivative Works, in at least one
+    of the following places: within a NOTICE text file distributed
+    as part of the Derivative Works; within the Source form or
+    documentation, if provided along with the Derivative Works; or,
+    within a display generated by the Derivative Works, if and
+    wherever such third-party notices normally appear. The contents
+    of the NOTICE file are for informational purposes only and
+    do not modify the License. You may add Your own attribution
+    notices within Derivative Works that You distribute, alongside
+    or as an addendum to the NOTICE text from the Work, provided
+    that such additional attribution notices cannot be construed
+    as modifying the License.
+
+    You may add Your own copyright statement to Your modifications and
+    may provide additional or different license terms and conditions
+    for use, reproduction, or distribution of Your modifications, or
+    for any such Derivative Works as a whole, provided Your use,
+    reproduction, and distribution of the Work otherwise complies with
+    the conditions stated in this License.
+
+5.  Submission of Contributions. Unless You explicitly state otherwise,
+    any Contribution intentionally submitted for inclusion in the Work
+    by You to the Licensor shall be under the terms and conditions of
+    this License, without any additional terms or conditions.
+    Notwithstanding the above, nothing herein shall supersede or modify
+    the terms of any separate license agreement you may have executed
+    with Licensor regarding such Contributions.
+
+6.  Trademarks. This License does not grant permission to use the trade
+    names, trademarks, service marks, or product names of the Licensor,
+    except as required for reasonable and customary use in describing the
+    origin of the Work and reproducing the content of the NOTICE file.
+
+7.  Disclaimer of Warranty. Unless required by applicable law or
+    agreed to in writing, Licensor provides the Work (and each
+    Contributor provides its Contributions) on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+    implied, including, without limitation, any warranties or conditions
+    of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+    PARTICULAR PURPOSE. You are solely responsible for determining the
+    appropriateness of using or redistributing the Work and assume any
+    risks associated with Your exercise of permissions under this License.
+
+8.  Limitation of Liability. In no event and under no legal theory,
+    whether in tort (including negligence), contract, or otherwise,
+    unless required by applicable law (such as deliberate and grossly
+    negligent acts) or agreed to in writing, shall any Contributor be
+    liable to You for damages, including any direct, indirect, special,
+    incidental, or consequential damages of any character arising as a
+    result of this License or out of the use or inability to use the
+    Work (including but not limited to damages for loss of goodwill,
+    work stoppage, computer failure or malfunction, or any and all
+    other commercial damages or losses), even if such Contributor
+    has been advised of the possibility of such damages.
+
+9.  Accepting Warranty or Additional Liability. While redistributing
+    the Work or Derivative Works thereof, You may choose to offer,
+    and charge a fee for, acceptance of support, warranty, indemnity,
+    or other liability obligations and/or rights consistent with this
+    License. However, in accepting such obligations, You may act only
+    on Your own behalf and on Your sole responsibility, not on behalf
+    of any other Contributor, and only if You agree to indemnify,
+    defend, and hold each Contributor harmless for any liability
+    incurred by, or claims asserted against, such Contributor by reason
+    of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+Copyright 2025 OpenAI
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

NOTICE

@@ -0,0 +1,19 @@
+JevCodeLocator-0.6B
+Copyright 2026 the JevCodeLocator authors
+
+This product includes model weights derived from:
+
+    Qwen/Qwen3-0.6B
+    Copyright 2025 Alibaba Cloud
+    Licensed under the Apache License, Version 2.0
+    https://huggingface.co/Qwen/Qwen3-0.6B
+
+The inference code in this repository (modeling_jev.py, inference_example.py,
+patch_tokenizer_for_transformers4.py) is original work, licensed under the
+Apache License, Version 2.0.
+
+Training data was produced programmatically (tree-sitter chunking, doc-comment /
+symbol / constant query generation, compact candidate summaries) from
+open-source repositories. See the section "Training data and licensing" in
+README.md for the repository list and their licenses. No source file from those
+repositories is redistributed here.

README.md

@@ -1,3 +1,177 @@
 ---
+language:
+  - en
+  - zh
+library_name: pytorch
 license: apache-2.0
+base_model: Qwen/Qwen3-0.6B
+tags:
+  - code-localization
+  - code-search
+  - code-retrieval
+  - reranker
+  - qwen3
+  - set-attention
+pipeline_tag: text-ranking
 ---
+
+# JevCodeLocator-0.6B
+
+A 0.6B **code localizer**. Given a repository and a natural-language query (an issue, a symptom, a
+symbol, an error message, or a described behavior), it reranks candidate code locations and returns a
+ranked list of `path:start-end` chunks together with a file-level ranking.
+
+The model is a **reranker**, not a retriever: a cheap sparse retriever (BM25 plus symbol/path channels)
+produces a shortlist of K candidates and the model scores all of them in a single forward pass. There
+are no embeddings and no vector index. On an entry-level 6 GB GPU it runs in roughly 0.4 s at K=8 and
+0.9 s at K=16, which makes it usable as a local first-stage localizer for coding agents.
+
+## Architecture
+
+```
+state (task / repo / candidate_count / query)  +  candidate texts ("path:start-end: summary")
+        │
+        ├─ one leaf-token path per candidate  ──►  Qwen3-0.6B backbone
+        │                                          (last hidden state at the EOS of each path)
+        ├─ LayerNorm → Linear(d, 1)                    per-candidate logit z_i
+        └─ set-attention over the K candidates:
+             u     = Linear(d + 1, 128)([h_i ; log K])
+             mixed = MultiheadAttention(u, u, u)       (key-padding mask for empty slots)
+             z_i  += Linear(128, 1)(tanh(u + mixed))
+        │
+        └─ softmax over the K candidates
+```
+
+The scalar head alone can rank candidates independently; the zero-initialized set-attention
+correction is what makes a candidate's score depend on the **whole shortlist** (for example,
+suppressing a cluster of look-alike test files in favour of the implementation file).
+
+## Training
+
+Three consecutive fine-tuning stages on top of `Qwen/Qwen3-0.6B`:
+
+| stage | training data | steps | notes |
+|---|---|---:|---|
+| 1 | Python programmatic queries, K≤16, file-level auxiliary loss λ=0.30 | 600 | base localizer |
+| 2 | + Go / TypeScript-JavaScript / C++ / Rust / C# / PHP programmatic queries; 12,274 rows, K=16 BM25 pools with the gold candidate forced in | 600 | doc-comment queries oversampled 2×, backbone LR 1e-5 |
+| 3 | continuation of stage 2 on 18,688 rows | 600 | final checkpoint |
+
+- Objective: cross-entropy against the gold candidate (`gold_distribution`) plus the file-level
+  auxiliary loss (weight 0.30).
+- Candidate text is identical at training and serving time: a compact summary consisting of the path,
+  the signature and the first lines of the chunk.
+- A Python replay set (2,286 rows) is kept in every round so that Python behaviour does not regress.
+- Training queries are generated programmatically from doc comments, symbol names and constants.
+  No human-written query set was used for training.
+
+## Evaluation
+
+Frozen held-out pools: 200 queries per repository, K=8, gold candidate forced into the pool.
+`file@1` is the share of queries whose gold **file** is ranked first. The baseline is the same model
+before the multilingual stages (stage 1 only).
+
+| evaluation set | language | baseline | **this model** |
+|---|---|---:|---:|
+| service repository, ~1.3k files | Go | 0.700 | **0.945** |
+| ML framework core, ~2k chunks | C++ | 0.730 | **0.870** |
+| agent CLI, ~20k chunks | Rust | 0.900 | 0.925 |
+| API service, ~5k chunks | TypeScript | 0.795 | 0.855 |
+| language server, ~600 chunks | C++ | 0.974 | 0.989 |
+| held-out dev set, 460 queries | Python | 0.839 | 0.843 |
+
+Restricted to semantic (doc-comment) queries, `file@1`: Go 0.524 → **0.913**, C++ 0.629 → 0.814,
+TypeScript 0.667 → 0.758, Rust 0.833 → 0.875.
+
+On a set of five hand-checked, real reverse-proxy questions about a large Go service (K=16, FP32):
+**gold-file recall@1 = 1.000, recall@3 = 1.000, recall@5 = 1.000.** Latency on an entry-level 6 GB
+GPU: 0.40–0.44 s (K=8), ~0.83 s (K=16).
+
+## Usage
+
+```bash
+pip install torch transformers safetensors
+python inference_example.py           # runs on this directory
+```
+
+```python
+import torch
+
+from modeling_jev import load_jev_model, score_candidates
+
+model, tok = load_jev_model(".", device="cuda" if torch.cuda.is_available() else "cpu",
+                            dtype=torch.float32)
+
+candidates = {
+    "src/auth/session.py:41-88": "def create_session(user, password) | validates credentials ...",
+    "src/db/pool.py:12-60": "def connect(dsn, max_size) | opens the database pool ...",
+}
+
+for key, prob in score_candidates(model, tok, "where are credentials validated?", candidates):
+    print(f"{prob:.4f}  {key}")
+```
+
+`score_candidates` returns the candidates ranked by probability. Summing the probabilities of all
+candidates that belong to the same file gives the file-level ranking.
+
+## Files
+
+| file | contents |
+|---|---|
+| `model.safetensors` | full model: 322 tensors (`backbone.*` plus the decision head); sha256 `fb06ba6ac6d30c395912ebd156723a0739afdd260f9b44f85c56b074c68b8a98` |
+| `config.json` | model configuration (`architectures: JevDecisionModel`) and training hyperparameters |
+| `backbone_config/config.json` | Qwen3-0.6B backbone configuration |
+| `tokenizer/` | tokenizer files (`tokenizer_config.json` is in the `transformers>=5` format) |
+| `head.safetensors` | the 12 decision-head tensors only (803 KB), for custom runtimes |
+| `modeling_jev.py` | self-contained model class, loader and scoring helper |
+| `inference_example.py` | runnable reranking example |
+| `patch_tokenizer_for_transformers4.py` | converts `tokenizer_config.json` to the 4.x format |
+| `LICENSE`, `NOTICE` | Apache-2.0 and attribution |
+
+## Intended use and limitations
+
+- **Use it as a reranker.** It selects among the candidates it is given; if the retriever's shortlist
+  does not contain the gold file, the model cannot recover it. Pool recall is the ceiling.
+- Trained on programmatically generated queries (doc comments, symbols, constants) with a small amount
+  of issue-style data mixed in. Free-form questions work but are the hardest distribution.
+- Language coverage: Python, Go, TypeScript/JavaScript, C, C++, Rust, C#, PHP, Java, Ruby, Bash. The
+  chunker used to build candidates is tree-sitter based; other languages are untested.
+- Candidate text is a compact summary, so very long functions are effectively truncated.
+- Not instruction-tuned, not a chat model, no safety tuning. Do not use it to make decisions about
+  code it cannot see.
+
+## Training data and licensing
+
+The weights in this repository are released under **Apache-2.0**, the same license as the base model
+[Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) (Apache-2.0, Copyright 2025 Alibaba Cloud).
+The inference code here (`modeling_jev.py`, `inference_example.py`,
+`patch_tokenizer_for_transformers4.py`) is original work under the same license; see `LICENSE` and
+`NOTICE`.
+
+Training rows are generated programmatically: a tree-sitter chunker extracts code locations, queries
+are derived from doc comments, symbol names and constants, and the candidate text is a compact summary
+of the chunk. **No source file from any repository is redistributed in this repository.**
+
+Licenses of the source repositories used for training and evaluation data:
+
+| license | sources |
+|---|---|
+| MIT | 6 repositories (Go / TypeScript / Python tooling) |
+| Apache-2.0 | 3 repositories (C++ / Go / Rust) |
+| GPL-3.0 | 2 repositories |
+| LGPL-3.0 | 2 repositories |
+| AGPL-3.0 | 1 repository |
+| non-commercial custom license | 1 repository |
+| license not declared in the checkout used | 6 repositories |
+
+A detailed list of the source repositories is available on request.
+
+Notes:
+
+- The weights are not a redistribution of those repositories, and the model does not reproduce training
+  code verbatim: it only scores candidate locations supplied by the caller. Whether copyleft-licensed
+  training data extends to model weights is legally unsettled in most jurisdictions.
+- **Commercial use**: the training mix includes one AGPL-3.0 project and one non-commercial project. For
+  a commercially clean release, retrain without those two (and preferably without the GPL/LGPL
+  projects); the data pipeline is deterministic and reproducible.
+- You remain responsible for complying with the licenses of the code you run the model on.
+- This section is provenance information, not legal advice.

backbone_config/config.json

@@ -0,0 +1,63 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "float32",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 28,
+  "model_type": "qwen3",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "pad_token_id": null,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

config.json

@@ -0,0 +1,50 @@
+{
+  "architectures": [
+    "JevDecisionModel"
+  ],
+  "model_type": "jev-decision",
+  "base_model": "Qwen/Qwen3-0.6B",
+  "backbone_model_type": "qwen3",
+  "backbone_config_path": "backbone_config",
+  "head_weights": "head.safetensors",
+  "set_head": "attention",
+  "task": "closed-set choice over candidate code locations (path:start-end)",
+  "candidate_input": "state (task/repo/candidate_count/query) + per-candidate summary text",
+  "languages": [
+    "python",
+    "go",
+    "typescript",
+    "javascript",
+    "c",
+    "cpp",
+    "rust",
+    "csharp",
+    "php",
+    "java",
+    "ruby",
+    "bash"
+  ],
+  "max_length": 2048,
+  "parameter_count": 596250498,
+  "parameter_storage": "float32",
+  "forward_autocast": "bfloat16",
+  "objective": "gold_distribution",
+  "file_loss_weight": 0.3,
+  "file_target_mode": "primary",
+  "loss": "ce",
+  "steps": 600,
+  "batch_questions": 32,
+  "microbatch_questions": 4,
+  "max_microbatch_tokens": 32768,
+  "backbone_lr": 3e-06,
+  "head_lr": 5e-05,
+  "lr_warmup_frac": 0.05,
+  "lr_min_frac": 0.1,
+  "precision": "bf16",
+  "gradient_checkpointing": true,
+  "train_questions": 18688,
+  "seed": 41,
+  "set_head_hidden": 128,
+  "set_head_heads": 4,
+  "num_parameters": 596250498
+}

head.safetensors

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

inference_example.py

@@ -0,0 +1,51 @@
+"""Minimal end-to-end example for JevCodeLocator-0.6B.
+
+    python inference_example.py
+
+Shows the intended serving pattern: a cheap retriever (BM25 / symbol search)
+produces a shortlist of `path:start-end` candidates, and this model reranks them
+with one forward pass.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+from modeling_jev import load_jev_model, score_candidates
+
+# Directory containing model.safetensors + config.json (default: this folder).
+CHECKPOINT = sys.argv[1] if len(sys.argv) > 1 else str(Path(__file__).resolve().parent)
+
+# A shortlist as a retriever would produce it: key = path:start-end, value = compact summary.
+CANDIDATES = {
+    "src/auth/session.py:41-88": "def create_session(user, password) | validates credentials against the user store and returns a signed session token",
+    "src/db/pool.py:12-60": "def connect(dsn, max_size) | opens and configures the database connection pool",
+    "src/api/routes/login.py:18-35": "def login_handler(request) | parses the JSON body and calls create_session, mapping failures to HTTP 401",
+    "src/auth/tokens.py:7-29": "def sign(payload, key) | HMAC-signs a payload with the application secret",
+    "docs/auth.md:1-40": "Authentication overview and deployment notes",
+}
+
+QUERY = "where are user credentials validated before a session token is created?"
+
+
+def main() -> None:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, tokenizer = load_jev_model(CHECKPOINT, device=device, dtype=torch.float32)
+    ranked = score_candidates(model, tokenizer, QUERY, CANDIDATES, repo="example", max_length=2048)
+    print(f"query: {QUERY}\n")
+    for rank, (key, prob) in enumerate(ranked, 1):
+        print(f"{rank}. {prob:.4f}  {key}")
+
+    # File-level aggregation: sum the probabilities of all candidates in one file.
+    files: dict[str, float] = {}
+    for key, prob in ranked:
+        files[key.rsplit(":", 1)[0]] = files.get(key.rsplit(":", 1)[0], 0.0) + prob
+    print("\nfile ranking:")
+    for path, score in sorted(files.items(), key=lambda kv: -kv[1]):
+        print(f"  {score:.4f}  {path}")
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

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

modeling_jev.py

@@ -0,0 +1,161 @@
+"""Self-contained inference code for JevCodeLocator-0.6B.
+
+The model is a Qwen3-0.6B backbone followed by a small set-attention head that
+scores a *set* of candidate code locations (``path:start-end``) against a query:
+
+    backbone(last hidden state at EOS of each candidate path)
+      -> LayerNorm -> scalar head            (per-candidate logit)
+      -> set-attention over the K candidates (context-dependent correction)
+
+One forward pass scores every candidate at once; no embeddings or vector index
+are involved. Candidate text is the compact summary produced by the data
+pipeline (``path | signature + first lines``).
+
+Example
+-------
+>>> from modeling_jev import load_jev_model, score_candidates
+>>> model, tok = load_jev_model("JevCodeLocator-0.6B", device="cuda")
+>>> cands = {
+...     "src/auth.py:10-42": "def login(user, pwd) | validates credentials and returns a session token",
+...     "src/db.py:5-30": "def connect(dsn) | opens a database connection pool",
+... }
+>>> for key, prob in score_candidates(model, tok, "where are credentials validated?", cands):
+...     print(f"{prob:.4f}  {key}")
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict, List, Sequence, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+INSTRUCTIONS = ("Select the single code location (file path and line range) that best answers the query. "
+                "Each option is a candidate location; its text includes a compact signature/code summary. "
+                "Only use the state and the option texts.")
+
+
+class JevDecisionModel(nn.Module):
+    """Qwen3 backbone + LayerNorm/scalar head + set-attention over candidates."""
+
+    def __init__(self, backbone: nn.Module, set_head: str = "attention") -> None:
+        super().__init__()
+        self.backbone = backbone
+        hidden = backbone.config.hidden_size
+        self.norm = nn.LayerNorm(hidden)
+        self.scalar = nn.Linear(hidden, 1)
+        nn.init.normal_(self.scalar.weight, std=0.02)
+        nn.init.zeros_(self.scalar.bias)
+        self.set_head = set_head
+        if set_head == "attention":
+            self.set_project = nn.Linear(hidden + 1, 128)
+            self.set_attention = nn.MultiheadAttention(128, 4, dropout=0.0, batch_first=True)
+            self.set_output = nn.Linear(128, 1)
+            nn.init.zeros_(self.set_output.weight)
+            nn.init.zeros_(self.set_output.bias)
+
+    def forward(self, examples: Sequence[dict], pad_token: int):
+        """examples: [{leaf_tokens: [[int, ...], ...], candidate_ids: [...], type: "choice"}]."""
+        paths = [ids for ex in examples for ids in ex["leaf_tokens"]]
+        device = self.scalar.weight.device
+        lengths = torch.tensor([len(ids) for ids in paths], device=device)
+        width = int(lengths.max())
+        tokens = torch.full((len(paths), width), pad_token, dtype=torch.long, device=device)
+        for i, ids in enumerate(paths):
+            tokens[i, :len(ids)] = torch.tensor(ids, device=device)
+        attention = torch.arange(width, device=device)[None, :] < lengths[:, None]
+        hidden = self.backbone(input_ids=tokens, attention_mask=attention,
+                               use_cache=False).last_hidden_state
+        leaves = hidden[torch.arange(len(paths), device=device), lengths - 1]
+        kmax = max(len(ex["candidate_ids"]) for ex in examples)
+        h = leaves.new_zeros((len(examples), kmax, leaves.shape[-1]))
+        valid = torch.zeros((len(examples), kmax), dtype=torch.bool, device=device)
+        offset = 0
+        for i, ex in enumerate(examples):
+            n = len(ex["leaf_tokens"])
+            h[i, :n] = leaves[offset:offset + n]
+            valid[i, :len(ex["candidate_ids"])] = True
+            offset += n
+        h = self.norm(h)
+        z = self.scalar(h).squeeze(-1).float()
+        choice = torch.tensor([i for i, ex in enumerate(examples) if ex["type"] == "choice"], device=device)
+        if self.set_head == "attention" and len(choice):
+            log_k = valid[choice].sum(-1).float().log()[:, None, None].expand(-1, kmax, 1)
+            u = self.set_project(torch.cat([h[choice], log_k.to(h.dtype)], dim=-1))
+            mixed, _ = self.set_attention(u, u, u, key_padding_mask=~valid[choice], need_weights=False)
+            delta = self.set_output(torch.tanh(u + mixed)).squeeze(-1).float()
+            z = z.index_add(0, choice, delta)
+        out = []
+        for i, ex in enumerate(examples):
+            if ex["type"] == "boolean":
+                out.append(F.pad(torch.stack([z[i, 0] * 0, z[i, 0]]), (0, kmax - 2)))
+            else:
+                out.append(z[i])
+        return torch.stack(out).masked_fill(~valid, -1e9), valid
+
+
+def load_jev_model(checkpoint_dir: str | Path, device: str = "cuda", dtype=torch.float32):
+    """Load the backbone, head weights and tokenizer from a checkpoint directory."""
+    from safetensors.torch import load_file
+    from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+    root = Path(checkpoint_dir).expanduser().resolve()
+    cfg = json.loads((root / "config.json").read_text(encoding="utf-8"))
+    body_config = AutoConfig.from_pretrained(str(root / "backbone_config"), local_files_only=True)
+    backbone = AutoModel.from_config(body_config, attn_implementation="sdpa")
+    model = JevDecisionModel(backbone, set_head=cfg.get("set_head", "attention"))
+    weights = root / "model.safetensors"
+    if not weights.is_file():
+        weights = root / "best.safetensors"
+    missing, unexpected = model.load_state_dict(load_file(str(weights)), strict=False)
+    if any(k.startswith("backbone.") or k.startswith("norm.") or k.startswith("scalar.") for k in missing):
+        raise RuntimeError(f"checkpoint is missing head/backbone weights: {list(missing)[:5]}")
+    model = model.to(device=device, dtype=dtype).eval()
+    tokenizer = AutoTokenizer.from_pretrained(str(root / "tokenizer"), local_files_only=True)
+    return model, tokenizer
+
+
+def prepare_examples(payload: dict, tokenizer, max_length: int = 2048) -> List[dict]:
+    """Encode state + candidate texts into one leaf-token path per candidate."""
+    if tokenizer.eos_token_id is None:
+        raise ValueError("tokenizer has no eos_token_id")
+    examples = []
+    for row in payload["states"]:
+        for qid, q in row["questions"].items():
+            ids = list(q["criteria"])
+            texts = [f"{key}: {q['criteria'][key]}" for key in ids]
+            segments = [f"State:\n{row['state']}\n",
+                        f"Question type: choice\nQuestion:\n{q.get('instructions', INSTRUCTIONS)}\n"]
+            prefix = sum([tokenizer.encode(t, add_special_tokens=False) for t in segments], [])
+            leaves = [prefix + tokenizer.encode(f"Candidate:\n{t}\nDecision:", add_special_tokens=False)
+                      + [tokenizer.eos_token_id] for t in texts]
+            largest = max(map(len, leaves))
+            if largest > max_length:
+                raise ValueError(f"{row['id']}:{qid} candidate path is {largest} tokens > max_length={max_length}")
+            examples.append({"id": f"{row['id']}:{qid}", "state_id": row["id"], "qid": qid,
+                             "type": "choice", "candidate_ids": ids, "candidate_texts": texts,
+                             "leaf_tokens": leaves})
+    return examples
+
+
+@torch.inference_mode()
+def score_candidates(model, tokenizer, query: str, candidates: Dict[str, str],
+                     repo: str = "repo", max_length: int = 2048) -> List[Tuple[str, float]]:
+    """Return [(candidate_key, probability)] sorted by probability, best first.
+
+    ``candidates`` maps ``path:start-end`` -> compact summary text; order is preserved
+    in the input but the returned list is ranked by the model.
+    """
+    state = ("task: fast_context_code_location\n"
+             f"repo: {repo}\n"
+             f"candidate_count: {len(candidates)}\n"
+             "query:\n"
+             f"{query}\n")
+    row = {"id": "query", "state": state,
+           "questions": {"q_loc": {"type": "choice", "instructions": INSTRUCTIONS, "criteria": dict(candidates)}}}
+    ex = prepare_examples({"states": [row]}, tokenizer, max_length)[0]
+    logits, valid = model([ex], tokenizer.pad_token_id)
+    probs = logits[0, :len(ex["candidate_ids"])].float().softmax(-1).tolist()
+    return sorted(zip(ex["candidate_ids"], probs), key=lambda kv: -kv[1])

patch_tokenizer_for_transformers4.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+"""Convert the tokenizer config from the transformers>=5 format to the 4.x format.
+
+`transformers` 5 stores `extra_special_tokens` as a list of token strings; 4.x
+expects a mapping. Run once if you use transformers 4.x:
+
+    python patch_tokenizer_for_transformers4.py tokenizer/tokenizer_config.json
+"""
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+
+def main() -> None:
+    path = Path(sys.argv[1] if len(sys.argv) > 1 else "tokenizer/tokenizer_config.json")
+    cfg = json.loads(path.read_text(encoding="utf-8"))
+    extra = cfg.get("extra_special_tokens")
+    if isinstance(extra, list):
+        cfg["extra_special_tokens"] = {token: token for token in extra}
+        path.write_text(json.dumps(cfg, ensure_ascii=False, indent=2), encoding="utf-8")
+        print(f"patched {path}: extra_special_tokens list -> dict ({len(cfg['extra_special_tokens'])} entries)")
+    elif isinstance(extra, dict):
+        print(f"{path}: already in 4.x format")
+    else:
+        print(f"{path}: no extra_special_tokens field, nothing to do")
+
+
+if __name__ == "__main__":
+    main()

tokenizer/chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

tokenizer/tokenizer.json

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

tokenizer/tokenizer_config.json

@@ -0,0 +1,28 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}