| # SWE-Pruner ONNX (code-pruner) | |
| ONNX-converted version of [ayanami-kitasan/code-pruner](https://huggingface.co/ayanami-kitasan/code-pruner) for efficient CPU inference. | |
| ## Source | |
| - **Original Model**: [ayanami-kitasan/code-pruner](https://huggingface.co/ayanami-kitasan/code-pruner) (safetensors) | |
| - **Training Code**: [Ayanami1314/swe-pruner](https://github.com/Ayanami1314/swe-pruner) | |
| ## Architecture | |
| - **Backbone**: Qwen/Qwen3-Reranker-0.6B (28 layers, hidden=1024) | |
| - **Multi-layer Fusion**: Early (layer 7) + Middle (layer 14) + Final (layer 28) → fused_hidden=3072 | |
| - **Fusion**: 1-layer MultiheadAttention (8 heads) + LayerNorm | |
| - **Compression Head**: CRF-style (LayerNorm → Linear(3072,256) → GELU → Linear(256,2)) | |
| - **Output**: `token_scores` — sigmoid scores per token (0-1, higher = keep) | |
| ## Files | |
| | File | Description | | |
| |------|-------------| | |
| | `model.onnx` | Quantized ONNX model (uint8, ~607MB) | | |
| | `vocab.json` | BPE vocabulary (Qwen3 tokenizer) | | |
| | `merges.txt` | BPE merge rules | | |
| | `metadata.json` | Model metadata (token IDs, dimensions) | | |
| | `crf_params.npz` | CRF transition parameters (optional, for Viterbi decoding) | | |
| ## Usage | |
| ```python | |
| import onnxruntime as ort | |
| import numpy as np | |
| sess = ort.InferenceSession("model.onnx") | |
| input_ids = np.array([[...]], dtype=np.int64) # [1, seq_len] | |
| attention_mask = np.array([[...]], dtype=np.int64) # [1, seq_len] | |
| scores = sess.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})[0] | |
| # scores: [1, seq_len] float32, 0-1 range, higher = keep | |
| ``` | |
| ## Conversion Details | |
| - Exported with PyTorch 2.8 + transformers 4.57 | |
| - Opset version: 14 | |
| - Dynamic axes: batch and seq_len | |
| - Quantized: dynamic uint8 quantization | |
| - Causal mask patched for ONNX trace compatibility | |