ayjays132 commited on
Commit
c33608b
·
verified ·
1 Parent(s): 87768f1

Release Phillnet Mini Text-Vision v1.1.0

Browse files

Text-and-still-image model release with adaptive reasoning defaults, complete model card, examples, integrity manifest, and production deployment materials.

.dockerignore ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ .git
2
+ .gitignore
3
+ __pycache__
4
+ *.py[cod]
5
+ .env
6
+ .env.*
7
+ !.env.example
8
+ *.log
9
+ *.pid
10
+ *.zip
11
+ *.tar
12
+ *.tar.gz
13
+ .pytest_cache
14
+ .mypy_cache
15
+ .cache
.env.example ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copy this file to .env and replace the placeholder before exposing the API.
2
+ # Leave empty only for isolated local development.
3
+ PHILLNET_API_KEY=replace-with-a-long-random-secret
4
+
5
+ # Comma-separated browser origins allowed to call the API. Do not use * in production.
6
+ PHILLNET_CORS_ORIGINS=https://your-app.example.com
7
+
8
+ # The container's model directory. Keep the bundled default unless mounting a different model path.
9
+ MODEL_DIR=/app/model
10
+
11
+ # Request limits for base64 image input and JSON bodies.
12
+ PHILLNET_MAX_IMAGE_BYTES=10485760
13
+ PHILLNET_MAX_IMAGE_PIXELS=24000000
14
+ PHILLNET_MAX_REQUEST_BYTES=12582912
.gitattributes CHANGED
@@ -1,35 +1,8 @@
1
- *.7z filter=lfs diff=lfs merge=lfs -text
2
- *.arrow filter=lfs diff=lfs merge=lfs -text
3
- *.bin filter=lfs diff=lfs merge=lfs -text
4
- *.bz2 filter=lfs diff=lfs merge=lfs -text
5
- *.ckpt filter=lfs diff=lfs merge=lfs -text
6
- *.ftz filter=lfs diff=lfs merge=lfs -text
7
- *.gz filter=lfs diff=lfs merge=lfs -text
8
- *.h5 filter=lfs diff=lfs merge=lfs -text
9
- *.joblib filter=lfs diff=lfs merge=lfs -text
10
- *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
- *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
- *.model filter=lfs diff=lfs merge=lfs -text
13
- *.msgpack filter=lfs diff=lfs merge=lfs -text
14
- *.npy filter=lfs diff=lfs merge=lfs -text
15
- *.npz filter=lfs diff=lfs merge=lfs -text
16
- *.onnx filter=lfs diff=lfs merge=lfs -text
17
- *.ot filter=lfs diff=lfs merge=lfs -text
18
- *.parquet filter=lfs diff=lfs merge=lfs -text
19
- *.pb filter=lfs diff=lfs merge=lfs -text
20
- *.pickle filter=lfs diff=lfs merge=lfs -text
21
- *.pkl filter=lfs diff=lfs merge=lfs -text
22
- *.pt filter=lfs diff=lfs merge=lfs -text
23
- *.pth filter=lfs diff=lfs merge=lfs -text
24
- *.rar filter=lfs diff=lfs merge=lfs -text
25
  *.safetensors filter=lfs diff=lfs merge=lfs -text
26
- saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
- *.tar.* filter=lfs diff=lfs merge=lfs -text
28
- *.tar filter=lfs diff=lfs merge=lfs -text
29
- *.tflite filter=lfs diff=lfs merge=lfs -text
30
- *.tgz filter=lfs diff=lfs merge=lfs -text
31
- *.wasm filter=lfs diff=lfs merge=lfs -text
32
- *.xz filter=lfs diff=lfs merge=lfs -text
33
- *.zip filter=lfs diff=lfs merge=lfs -text
34
- *.zst filter=lfs diff=lfs merge=lfs -text
35
- *tfevents* filter=lfs diff=lfs merge=lfs -text
 
1
+ # Hugging Face Git LFS artifacts
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
  *.safetensors filter=lfs diff=lfs merge=lfs -text
3
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
4
+
5
+ # Keep Markdown, JSON configuration, HTML, and PNG examples as regular Git files
6
+ # for direct review in the repository viewer.
7
+ examples/direct-horizon-tasks.png filter=lfs diff=lfs merge=lfs -text
8
+ examples/medium-solara-dashboard.png filter=lfs diff=lfs merge=lfs -text
 
 
 
 
.gitignore ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Python bytecode and caches
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.so
5
+ .pytest_cache/
6
+ .mypy_cache/
7
+ .ruff_cache/
8
+
9
+ # Local environments and secrets
10
+ .venv/
11
+ venv/
12
+ .env
13
+ .env.*
14
+ !.env.example
15
+
16
+ # Local runtime data and logs
17
+ *.log
18
+ *.pid
19
+ .cache/
20
+ .tmp/
21
+
22
+ # Release archives are created outside the upload repository
23
+ *.zip
24
+ *.tar
25
+ *.tar.gz
26
+
27
+ # Operating-system files
28
+ .DS_Store
29
+ Thumbs.db
ADAPTIVE_DEFAULTS_REPORT.md ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Adaptive Reasoning and 8K Output Configuration
2
+
3
+ **Status:** Corrected and validated on 21 August 2026.
4
+
5
+ ## What was wrong with the prior HTML test
6
+
7
+ The model checkpoint already contained adaptive-reasoning features, but the evaluation caller overrode them. The prior benchmark explicitly passed `reasoning_effort="direct"`, a small manual recurrent budget, and `max_new_tokens=180`. Those three caller arguments bypassed the model's normal default reasoning profile and stopped every HTML response inside its opening CSS.
8
+
9
+ The issue was therefore not a missing SDXL removal or a tokenizer issue. It was a **generation-policy bottleneck in the testing and serving layer**.
10
+
11
+ ## Persisted defaults now enabled
12
+
13
+ The package configuration now uses the following defaults:
14
+
15
+ | Configuration field | Value | Effect |
16
+ |---|---|---|
17
+ | `adaptive_reasoning` | `true` | Enables adaptive recurrent reasoning behavior. |
18
+ | `default_adaptive_capsule_budget` | `true` | Allows adaptive reasoning capsules. |
19
+ | `default_adaptive_reasoning_trace` | `true` | Enables the private adaptive reasoning trace policy. |
20
+ | `default_private_reasoning_generation` | `true` | Makes `model.generate()` enter private reasoning when the caller does not override it. |
21
+ | `reasoning_token_budget_policy` | `adaptive_context` | Lets the private reasoning phase expand up to the real active-context boundary and stop naturally on `</think>`. |
22
+ | `default_reasoning_effort` | **`max`** | Makes maximum adaptive reasoning the package default. |
23
+ | `reasoning_effort_token_budgets.max` | `8192` | Retains the model's maximum reasoning profile. |
24
+ | `answer_token_budget_policy` | **`fixed`** | Prevents the visible-answer phase from expanding to the 1,048,576-token logical context. |
25
+ | `generation_default_max_answer_tokens` | **`8192`** | Gives visible responses an 8,192-token maximum by default. |
26
+
27
+ > “Infinite” reasoning is not physically literal: the model uses an adaptive policy bounded by the active 32,768-token cache window, the reserved answer allocation, and its end-of-thinking token. This is the correct safety and memory boundary. The model decides when to close its private reasoning phase; it is no longer artificially restricted by the former 180-token benchmark cap.
28
+
29
+ ## Serving defaults
30
+
31
+ `server.py` now defaults to:
32
+
33
+ ```python
34
+ max_tokens = 8192
35
+ reasoning_effort = "max"
36
+ ```
37
+
38
+ Callers can still request a lower `max_tokens` or reasoning effort when latency is more important than depth. The default path prioritizes the user-requested deep adaptive behavior.
39
+
40
+ ## Full uncapped adaptive HTML validation
41
+
42
+ The validation runner deliberately omitted `max_new_tokens`, `max_answer_tokens`, `max_reasoning_tokens`, and `reasoning_budget`. Consequently, the model used only persisted configuration defaults.
43
+
44
+ | Metric | Result |
45
+ |---|---:|
46
+ | Prompt tokens | 107 |
47
+ | Private-reasoning ceiling calculated from active context | 22,356 tokens |
48
+ | Private reasoning actually used before natural close | 2,944 tokens |
49
+ | Visible-answer ceiling | 8,192 tokens |
50
+ | Visible answer actually generated | 2,991 tokens |
51
+ | Total runtime on sandbox CPU | 921.415 seconds |
52
+ | Complete HTML document detected | Yes |
53
+ | HTML character count | 10,472 |
54
+
55
+ The output contains one complete document with balanced CSS braces, closing `style`, `head`, `body`, `script`, and `html` tags, a responsive mobile-toggle element, the `toggleMenu()` handler, a navigation container, two calls to action, and three feature cards. It was rendered locally in Chromium as a usable Orbit Notes landing page.
56
+
57
+ ## Important policy distinction
58
+
59
+ The first attempted configuration retained `answer_token_budget_policy: "adaptive_context"`. With a 1,048,576-token logical context, an omitted answer cap caused the private generator to allocate over one million visible-answer tokens after its reasoning phase. That was not an 8K maximum and exhausted the constrained evaluation process. The configuration was corrected to `answer_token_budget_policy: "fixed"` while preserving adaptive private reasoning. This gives the intended combination:
60
+
61
+ 1. **Adaptive deep private reasoning** that can self-terminate naturally before the active context boundary.
62
+ 2. **A bounded 8,192-token visible response** appropriate for practical serving and one-shot code generation.
63
+
64
+ ## Reproduction
65
+
66
+ Run the full adaptive validation with:
67
+
68
+ ```bash
69
+ python3 /home/ubuntu/phillnet_work/run_uncapped_adaptive_html_oneshot.py
70
+ ```
71
+
72
+ The artifacts are in `uncapped_adaptive_html_result/`:
73
+
74
+ | File | Description |
75
+ |---|---|
76
+ | `oneshot.html` | Complete generated Orbit Notes page. |
77
+ | `raw.txt` | Exact decoded model completion. |
78
+ | `summary.json` | Loaded defaults, completion totals, and progress events. |
79
+ | `progress.jsonl` | Privacy-safe phase/count/budget progress trail. |
COHERENCE_REPAIR_REPORT.md ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Coherence and Spatial-Input Repair Report
2
+
3
+ **Status:** Corrected and validated on 21 August 2026.
4
+
5
+ ## Root cause
6
+
7
+ The initial lean build retained the correct language-and-vision checkpoint but introduced a **tokenizer contract error** in the new image processor. It instantiated `DendroByteTokenizer`, a 259-token byte tokenizer, while the retained checkpoint declares a 248,320-token Qwen-derived vocabulary and includes a native `Qwen2Tokenizer` artifact. The mismatch sent byte-token IDs to embeddings trained for Qwen token IDs, causing malformed decoded output such as `�` even though the model could load and complete numerical forward passes.
8
+
9
+ A second prompt-framing issue caused the initial service to construct literal `\\n` character pairs rather than newline separators. That altered the chat token sequence and made the HTTP endpoint differ from the native Qwen chat template. A third serving issue passed a near-zero `temperature` to the custom generation loop even with sampling disabled, which altered its deterministic behavior.
10
+
11
+ ## Corrections
12
+
13
+ | Area | Previous behavior | Corrected behavior |
14
+ |---|---|---|
15
+ | Language tokenizer | Custom 259-token byte tokenizer | Native `Qwen2Tokenizer` loaded locally from `tokenizer.json` with remote code disabled for tokenizer loading |
16
+ | Chat framing | Literal backslash-n characters in the processor prompt | Native Qwen-compatible `<|im_start|>`, real newline, `<|im_end|>` sequence |
17
+ | Image placement | Image features appended after the assembled text | Image placeholder sequence inserted exactly where each image content item occurs |
18
+ | Spatial metadata | Not independently tested | `image_grid_thw`, image placeholder count, and assistant-prefix ordering tested |
19
+ | Deterministic service decoding | Passed `temperature=1e-5` while `do_sample=False` | Omits temperature during deterministic decoding, preserving native argmax behavior |
20
+ | Service process | Older detached process remained bound to port 8000 | Stale process removed and the corrected service bound cleanly to port 8000 |
21
+
22
+ ## Coherence validation
23
+
24
+ All text completions below use deterministic direct-mode generation and are taken from the corrected processor path. The wrapped processor now produces the same token count and token IDs as the native Qwen chat template.
25
+
26
+ | Prompt | Corrected completion | Result |
27
+ |---|---|---|
28
+ | `Complete this phrase with one word: The capital of France is` | `Paris` | Correct |
29
+ | `What is 2 plus 2? Answer with only the number.` | `4` | Correct |
30
+ | `Reply with exactly one word: blue` | `blue` | Correct |
31
+ | Same last prompt through live `/v1/chat/completions` | `blue` | Correct |
32
+
33
+ ## Vision and spatial validation
34
+
35
+ The processor prepares a still image as two temporal frames, 16×16 patches, and a Qwen-compatible `image_grid_thw`. It inserts one `<|image_pad|>` token per spatially merged vision feature and places the vision tokens before the assistant generation prefix.
36
+
37
+ | Test | Input | Output | Result |
38
+ |---|---|---|---|
39
+ | Color grounding | Solid red image | `red` | Correct |
40
+ | Color grounding | Solid lime image | `green` | Correct |
41
+ | Color grounding | Solid blue image | `blue` | Correct |
42
+ | Color grounding | Solid white image | `white` | Correct |
43
+ | Color grounding | Solid black image | `black` | Correct |
44
+ | Spatial grounding | 128×64 scene: red left half, blue right half; ask left color | `red` | Correct |
45
+ | Spatial grounding | Same scene; ask right color | `blue` | Correct |
46
+
47
+ For the two-region spatial scene, the processor emitted `image_grid_thw = [[1, 4, 8]]` and eight visual placeholder tokens, matching the vision tower’s spatial merge factor of two. The direct test verified that the rightmost image placeholder precedes the assistant prefix, so vision embeddings are injected into the intended context position.
48
+
49
+ ## Live API status
50
+
51
+ The temporary endpoint remains:
52
+
53
+ ```text
54
+ https://8000-iq1jte46c8ls1hzprn1yn-9c029be5.us5.manus.computer/health
55
+ ```
56
+
57
+ The live service now uses the corrected processor and deterministic generation path. For a dark navy image, it answered `black`; the dedicated visual probe shows correct output for pure blue (`blue`) and for five high-contrast solid colors. The navy result should be interpreted as a dark-color ambiguity rather than a failure of image-token placement.
58
+
59
+ ## Reproduction tests
60
+
61
+ The parent workspace contains these scripts:
62
+
63
+ | Script | Purpose |
64
+ |---|---|
65
+ | `diagnose_tokenizer_alignment.py` | Demonstrates the original byte-vocabulary mismatch and the native tokenizer path. |
66
+ | `test_text_coherence.py` | Compares native Qwen chat formatting against the corrected processor on three deterministic prompts. |
67
+ | `test_vision_coherence.py` | Verifies the still-image processor contract and a decoded response. |
68
+ | `test_vision_color_grounding.py` | Verifies that five distinct solid-color images produce distinct correct answers. |
69
+ | `test_vision_spatial_grounding.py` | Verifies left/right spatial grounding on a two-region image. |
70
+
71
+ The model has now been demonstrated to provide coherent text completion, still-image color grounding, and basic left/right spatial grounding. These are smoke tests, not a comprehensive capability benchmark.
Dockerfile ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.12-slim
2
+
3
+ ENV PYTHONDONTWRITEBYTECODE=1 \
4
+ PYTHONUNBUFFERED=1 \
5
+ PIP_NO_CACHE_DIR=1 \
6
+ MODEL_DIR=/app/model
7
+
8
+ WORKDIR /app/model
9
+
10
+ COPY requirements.txt requirements-server.txt ./
11
+ RUN pip install --no-cache-dir -r requirements.txt -r requirements-server.txt \
12
+ && useradd --create-home --uid 10001 --shell /usr/sbin/nologin modeluser
13
+
14
+ COPY --chown=modeluser:modeluser . ./
15
+
16
+ USER modeluser
17
+ EXPOSE 8000
18
+
19
+ HEALTHCHECK --interval=30s --timeout=5s --start-period=120s --retries=3 \
20
+ CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/ready', timeout=3)"
21
+
22
+ CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1", "--proxy-headers", "--forwarded-allow-ips", "*"]
HF_UPLOAD.md ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Upload Phillnet Mini Text-Vision to Hugging Face
2
+
3
+ This directory is the **complete upload root**. Upload the directory itself, not a nested archive, so the model card, `config.json`, custom model code, tokenizer, checkpoint, examples, and deployment materials remain at repository root.
4
+
5
+ ## Recommended CLI workflow
6
+
7
+ Install the current Hub client, authenticate with a token that has write access, create the model repository, then upload this directory. The official `hf upload` workflow streams folders, resumes interrupted uploads, and honours the root `.gitignore` file. [1]
8
+
9
+ ```bash
10
+ python -m pip install --upgrade huggingface_hub
11
+ hf auth login
12
+
13
+ # Choose your own account or organization namespace.
14
+ hf repo create YOUR_NAMESPACE/Phillnet-Mini-Text-Vision --type model
15
+
16
+ # Run from the parent directory that contains Phillnet-Mini-Text-Vision/.
17
+ export HF_XET_HIGH_PERFORMANCE=1
18
+ hf upload YOUR_NAMESPACE/Phillnet-Mini-Text-Vision \
19
+ ./Phillnet-Mini-Text-Vision \
20
+ . \
21
+ --commit-message "Release Phillnet Mini Text-Vision v1.1.0"
22
+ ```
23
+
24
+ The upload command can be rerun safely if a network interruption occurs; committed files are skipped and existing content is deduplicated by the Hub upload pipeline. [1]
25
+
26
+ ## Pre-upload checklist
27
+
28
+ | Item | Expected state |
29
+ |---|---|
30
+ | `README.md` | Model card with Hugging Face metadata and embedded local example gallery. |
31
+ | `model.safetensors` | Full retained language-and-vision checkpoint, tracked through Git LFS policy. |
32
+ | `config.json`, tokenizer files, and custom Python modules | Present at repository root for `trust_remote_code=True` loading. |
33
+ | `examples/` | Three complete HTML outputs and PNG previews. |
34
+ | `LICENSE` | Apache License 2.0 text. |
35
+ | `.gitattributes` | Tracks `*.safetensors` and `tokenizer.json` as large artifacts. |
36
+ | `.gitignore` | Excludes caches, secrets, local logs, and archives. |
37
+ | `Dockerfile`, `docker-compose.yml`, `.env.example` | Included for controlled self-hosted API deployment. |
38
+ | `RELEASE_MANIFEST.json` | Integrity and inventory record. |
39
+
40
+ ## Verify after upload
41
+
42
+ Open `https://huggingface.co/YOUR_NAMESPACE/Phillnet-Mini-Text-Vision` and verify that the README renders its three images, the Files tab contains the checkpoint and configuration, and the model card identifies this as a custom-code model requiring `trust_remote_code=True`.
43
+
44
+ Then test from a clean environment:
45
+
46
+ ```python
47
+ from transformers import AutoModelForCausalLM, AutoProcessor
48
+
49
+ repo = "YOUR_NAMESPACE/Phillnet-Mini-Text-Vision"
50
+ processor = AutoProcessor.from_pretrained(repo, trust_remote_code=True)
51
+ model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
52
+ ```
53
+
54
+ ## Git alternative
55
+
56
+ The Hub also supports Git-based model repositories. The included `.gitattributes` provides the necessary large-file rules, but the CLI workflow above is recommended for this full folder because it handles resumed large-file transfers automatically. [1] [2]
57
+
58
+ ## References
59
+
60
+ [1]: https://huggingface.co/docs/huggingface_hub/en/guides/upload "Hugging Face Hub: Upload files to the Hub"
61
+ [2]: https://huggingface.co/docs/hub/en/models-uploading "Hugging Face Hub: Uploading models"
IMPLEMENTATION_REPORT.md ADDED
@@ -0,0 +1,240 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Phillnet Mini Text-Vision
2
+
3
+ **Build status:** Complete and validated on 21 August 2026.
4
+
5
+ This package is a deliberately lean derivative of `ayjays132/Phillnet-Mini-Omni-Max`. It retains **text generation** and **still-image understanding** only. The upstream model card identifies the checkpoint as a custom-code Transformers model and describes a separate SDXL image-generation route; this build removes that route completely rather than merely disabling it at runtime. [1]
6
+
7
+ > **Scope lock:** This package supports only language generation and visual question answering over still images. It contains no SDXL U-Net, VAE, SDXL text encoders, diffusion scheduler, text-to-image API, text-to-video API, tool runtime, agent runtime, audio endpoint, or video endpoint.
8
+
9
+ ## What was downloaded and what was removed
10
+
11
+ The full original Hugging Face snapshot was mirrored before transformation. Its source inventory contained 99 files, including 18 files in the `phillnet3_sdxl/` subtree. The final lean package contains 48 files and is 1,785,035,602 bytes on disk, compared with 10,403,517,915 bytes for the complete downloaded snapshot.
12
+
13
+ | Artifact | Original snapshot | Final text-vision package | Result |
14
+ |---|---:|---:|---|
15
+ | Total package size | 10,403,517,915 bytes | 1,785,035,602 bytes | **82.84% reduction** |
16
+ | SDXL subtree | 8,608,200,761 bytes across 18 files | 0 bytes | **Removed** |
17
+ | Core `model.safetensors` | 1,763,655,304 bytes | 1,763,655,304 bytes | **Retained byte-for-byte** |
18
+ | Package files | 300 files in downloaded snapshot | 48 files | Runtime surface reduced |
19
+
20
+ The retained core checkpoint has SHA-256:
21
+
22
+ ```text
23
+ f1a913f99f8ce921c1aa79982a09eaee6744448766f09a4756751e2d4b9342fc
24
+ ```
25
+
26
+ That hash is identical before and after the stripping process. Therefore, the language and visual-encoder weights that reside in `model.safetensors` were not rewritten or quantized.
27
+
28
+ ## Exact removal policy
29
+
30
+ The SDXL assets were not left as optional downloads. They were physically excluded from the final package together with the loading paths that referenced them. This prevents the deployment from downloading, initializing, or accidentally exposing a diffusion backend.
31
+
32
+ | Removed area | Files or behavior removed | Why it is absent |
33
+ |---|---|---|
34
+ | SDXL checkpoint | `phillnet3_sdxl/`, including two `phillnet-image-*.safetensors` shards | These shards contain the diffusion donor banks and were the primary storage cost. |
35
+ | SDXL runtime | `phillnet3_sdxl.py`, `synthesis.py`, `phillnet3_bridge.py`, `generation_acceleration.py` | They create or accelerate the SDXL U-Net, VAE, CLIP encoders, scheduler, and image/video synthesis routes. |
36
+ | Generation APIs | `generate_image`, `generate_image_prompt`, `generate_video`, `generate_video_prompt`, `synthesis_losses` | These methods were removed from `DendroForCausalLM`, not merely hidden. |
37
+ | Diffusion configuration | SDXL/diffusion-specific fields and validation from `config.json` and `configuration_dendro_omni.py` | A lean model should not advertise, validate, or accept inactive synthesis configuration. |
38
+ | Non-core runtimes | Agent, tool, orchestration, Smolagents, legacy showcase, cache, and image-generation files | They are outside the requested text-and-vision serving surface. |
39
+ | External processor declaration | Stale `Qwen3VLProcessor` declaration | Replaced by the bundled `DendroVisionProcessor` to make image input local and self-contained. |
40
+
41
+ The remaining configuration declares `model_capabilities: ["text-generation", "image-understanding"]` and `text_vision_only: true`.
42
+
43
+ ## What remains and why it works
44
+
45
+ The central retained checkpoint uses a packed single-source SafeTensors layout. It has one top-level tensor rather than many separately named tensors, while `transplant_manifest.json` maps portions of that source to language and visual components. The manifest contains **153 retained `model.visual.*` mappings**. This is distinct from the deleted SDXL subtree: those mappings implement visual encoding for image understanding, whereas SDXL supplied image synthesis.
46
+
47
+ The new `processing_dendro_omni.py` is the key compatibility layer. It loads text and still images locally, normalizes an image, duplicates a still frame to satisfy the visual tower’s temporal patch dimension, constructs visual patches, inserts exactly the matching number of image placeholder tokens, and emits the `pixel_values`, `image_grid_thw`, and `mm_token_type_ids` tensors needed by the retained visual transformer. The processor is registered in both `config.json` and `preprocessor_config.json`, so `AutoProcessor.from_pretrained(..., trust_remote_code=True)` resolves to `DendroVisionProcessor`.
48
+
49
+ | Supported path | Status | Implementation |
50
+ |---|---|---|
51
+ | Text completion | Enabled | `model.generate(...)` with `reasoning_effort="direct"` for low-latency serving. |
52
+ | Text forward pass | Enabled | Standard `DendroForCausalLM.forward(...)`. |
53
+ | Still-image understanding | Enabled | `DendroVisionProcessor` plus `model.answer_image(...)` or direct multimodal model calls. |
54
+ | SDXL image generation | Removed | No weights, no module, no API, no dependency. |
55
+ | Video generation | Removed | No API, no diffusion route, no endpoint. |
56
+ | Audio / tools / agents | Not exposed | Excluded from the lean HTTP service and removed where they were optional runtime layers. |
57
+
58
+ ## Verification evidence
59
+
60
+ Verification was intentionally separated into structural, load, forward, and serving checks. The checks prove that the package can load and execute the retained text and visual-input paths. They do **not** constitute a semantic-quality benchmark; output quality must be evaluated separately on representative user tasks.
61
+
62
+ | Check | Result | Evidence |
63
+ |---|---|---|
64
+ | Python syntax compilation | Passed | Lean runtime modules, processor, configuration, model, and service compiled. |
65
+ | SDXL filesystem scan | Passed | No SDXL-named source or weight artifact remains in the final package. |
66
+ | Synthesis-code scan | Passed | No `diffusers`, `StableDiffusion`, `DendroSharedDiffusion`, image/video-generation method, or SDXL adapter reference remains in Python runtime code. |
67
+ | Configuration deserialization | Passed | `DendroOmniConfig.from_pretrained(...)` loaded from the stripped directory. |
68
+ | Processor discovery | Passed | `AutoProcessor` resolves to `DendroVisionProcessor`. |
69
+ | Model loading | Passed | `AutoModelForCausalLM` instantiated `DendroForCausalLM` in BF16 on CPU. |
70
+ | Text forward pass | Passed | Finite logits with shape `[1, 27, 248320]`. |
71
+ | Image forward pass | Passed | Finite logits with shape `[1, 61, 248320]`. |
72
+ | Visual-input contract | Passed | One 32×32 test image produced one visual placeholder and four visual patches. |
73
+ | Direct text generation | Passed | Two output tokens completed in 2.224 seconds in the sandbox CPU test. |
74
+ | Local deployment health | Passed | `/health` returns exactly the two enabled capabilities and the disabled set. |
75
+ | Local text endpoint | Passed | `/v1/chat/completions` completed a direct-mode text request. |
76
+ | Local image endpoint | Passed | `/v1/chat/completions` accepted a valid base64 still image plus text question. |
77
+ | Public temporary health check | Passed | Temporary proxied endpoint returned the expected health payload. |
78
+
79
+ ## Local use
80
+
81
+ Install the model dependencies from this directory, then load it using the local custom-code implementation:
82
+
83
+ ```bash
84
+ cd Phillnet-Mini-Text-Vision
85
+ pip install -r requirements.txt
86
+ ```
87
+
88
+ ```python
89
+ import torch
90
+ from transformers import AutoModelForCausalLM, AutoProcessor
91
+
92
+ model_dir = "./Phillnet-Mini-Text-Vision"
93
+ processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)
94
+ model = AutoModelForCausalLM.from_pretrained(
95
+ model_dir,
96
+ trust_remote_code=True,
97
+ dtype=torch.bfloat16,
98
+ low_cpu_mem_usage=True,
99
+ ).eval()
100
+
101
+ encoded = processor("Reply with one word: blue", return_tensors="pt")
102
+ output = model.generate(
103
+ **encoded,
104
+ max_new_tokens=32,
105
+ reasoning_effort="direct",
106
+ do_sample=False,
107
+ )
108
+ print(processor.tokenizer.decode(output[0], skip_special_tokens=True))
109
+ ```
110
+
111
+ For still-image questions, pass OpenAI-style multimodal message content to the processor:
112
+
113
+ ```python
114
+ from PIL import Image
115
+
116
+ encoded = processor.apply_chat_template(
117
+ [{
118
+ "role": "user",
119
+ "content": [
120
+ {"type": "image", "image": Image.open("example.png").convert("RGB")},
121
+ {"type": "text", "text": "Describe the visible objects."},
122
+ ],
123
+ }],
124
+ tokenize=True,
125
+ add_generation_prompt=True,
126
+ return_dict=True,
127
+ return_tensors="pt",
128
+ )
129
+
130
+ answer = model.generate(
131
+ **encoded,
132
+ max_new_tokens=128,
133
+ reasoning_effort="direct",
134
+ do_sample=False,
135
+ )
136
+ print(processor.tokenizer.decode(answer[0], skip_special_tokens=True))
137
+ ```
138
+
139
+ ## Deployment
140
+
141
+ The package contains a minimal FastAPI service in `server.py` and a `Dockerfile`. The service exposes only `/health` and `/v1/chat/completions`.
142
+
143
+ ### Run directly
144
+
145
+ ```bash
146
+ cd Phillnet-Mini-Text-Vision
147
+ pip install -r requirements.txt -r requirements-server.txt
148
+ uvicorn server:app --host 0.0.0.0 --port 8000
149
+ ```
150
+
151
+ ### Run in a container
152
+
153
+ ```bash
154
+ docker build -t phillnet-mini-text-vision .
155
+ docker run --rm -p 8000:8000 phillnet-mini-text-vision
156
+ ```
157
+
158
+ The package was verified on CPU using BF16 weights. Production capacity must account for the 1.76 GB checkpoint plus framework and request-memory overhead; test the target hardware under representative image sizes and concurrent-request volume before enabling unrestricted traffic.
159
+
160
+ ### Health check
161
+
162
+ ```bash
163
+ curl http://localhost:8000/health
164
+ ```
165
+
166
+ Expected response:
167
+
168
+ ```json
169
+ {
170
+ "status": "ok",
171
+ "capabilities": ["text-generation", "image-understanding"],
172
+ "disabled": ["image-generation", "video-generation", "audio", "tools", "agents"]
173
+ }
174
+ ```
175
+
176
+ ### Text request
177
+
178
+ ```bash
179
+ curl http://localhost:8000/v1/chat/completions \
180
+ -H 'Content-Type: application/json' \
181
+ --data '{
182
+ "messages": [{"role": "user", "content": "Give a one-sentence summary of caching."}],
183
+ "max_tokens": 96,
184
+ "reasoning_effort": "direct"
185
+ }'
186
+ ```
187
+
188
+ ### Still-image question
189
+
190
+ The endpoint accepts `image_base64`; remote image URLs are intentionally not fetched, avoiding a server-side request-forgery surface.
191
+
192
+ ```bash
193
+ IMAGE_B64=$(base64 -w0 example.png)
194
+ curl http://localhost:8000/v1/chat/completions \
195
+ -H 'Content-Type: application/json' \
196
+ --data "{
197
+ \"messages\": [{
198
+ \"role\": \"user\",
199
+ \"content\": [
200
+ {\"type\": \"image\", \"image_base64\": \"${IMAGE_B64}\"},
201
+ {\"type\": \"text\", \"text\": \"What is visible in this image?\"}
202
+ ]
203
+ }],
204
+ \"max_tokens\": 128,
205
+ \"reasoning_effort\": \"direct\"
206
+ }"
207
+ ```
208
+
209
+ ## Temporary deployment
210
+
211
+ A verified **temporary** endpoint is currently available at:
212
+
213
+ ```text
214
+ https://8000-iq1jte46c8ls1hzprn1yn-9c029be5.us5.manus.computer
215
+ ```
216
+
217
+ Its health endpoint is:
218
+
219
+ ```text
220
+ https://8000-iq1jte46c8ls1hzprn1yn-9c029be5.us5.manus.computer/health
221
+ ```
222
+
223
+ This URL is for validation only. It is backed by the session sandbox and will not provide durable production hosting. For a persistent deployment, use the included container definition on a host with enough RAM for the retained BF16 checkpoint and configure authentication, TLS termination, logging, rate limits, and request-size limits at the deployment boundary.
224
+
225
+ ## Reproducibility artifacts
226
+
227
+ The parent workspace contains the following non-runtime artifacts:
228
+
229
+ | File | Purpose |
230
+ |---|---|
231
+ | `build_text_vision_only.py` | Rebuilds the lean package from the fully downloaded original snapshot without executing remote model code. |
232
+ | `validate_text_vision_package.py` | Structural, configuration, processor, and checkpoint-inventory validation. |
233
+ | `load_text_vision_model.py` | Controlled model-load verification. |
234
+ | `run_text_vision_forward.py` | Text, image, and direct-generation forward smoke test. |
235
+ | `original_weight_sha256.txt` | Checksums captured from the complete original snapshot. |
236
+ | `SOURCE_INVENTORY.md` | Source repository inventory. |
237
+
238
+ ## References
239
+
240
+ [1]: https://huggingface.co/ayjays132/Phillnet-Mini-Omni-Max "ayjays132/Phillnet-Mini-Omni-Max model card"
PRODUCTION.md ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Production Operations Guide
2
+
3
+ Phillnet Mini Text-Vision is a **single-process, text-and-still-image inference service**. The included stack is designed for a controlled deployment boundary: one model process, one serialized generation at a time, local base64 image inputs only, optional API-key authentication, and explicit CORS origins.
4
+
5
+ ## Deployment topology
6
+
7
+ ```text
8
+ Trusted clients
9
+ │ HTTPS + API key
10
+
11
+ TLS reverse proxy / edge rate limiter
12
+ │ localhost HTTP
13
+
14
+ Docker Compose: Phillnet Mini Text-Vision
15
+ │ one Uvicorn worker + one loaded model
16
+
17
+ Text generation and still-image understanding
18
+ ```
19
+
20
+ Keep the model container bound to loopback by default. Terminate TLS, rate limit, and log access at a reverse proxy or managed ingress. The model container itself is intentionally not a public multi-tenant gateway.
21
+
22
+ ## First deployment
23
+
24
+ ```bash
25
+ cp .env.example .env
26
+ # Set PHILLNET_API_KEY to a long random secret.
27
+ # Set PHILLNET_CORS_ORIGINS to the exact browser application origin.
28
+ docker compose up --build -d
29
+
30
+ docker compose ps
31
+ curl http://127.0.0.1:8000/ready
32
+ ```
33
+
34
+ The service reports `ready: true` only after the processor and checkpoint are loaded.
35
+
36
+ ## Reverse-proxy example
37
+
38
+ The following Nginx location keeps the container private, applies a request body limit, and forwards the authorization header. Configure a valid TLS server block around it.
39
+
40
+ ```nginx
41
+ location / {
42
+ proxy_pass http://127.0.0.1:8000;
43
+ proxy_http_version 1.1;
44
+ proxy_set_header Host $host;
45
+ proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
46
+ proxy_set_header X-Forwarded-Proto $scheme;
47
+ proxy_set_header Authorization $http_authorization;
48
+ proxy_set_header X-API-Key $http_x_api_key;
49
+ client_max_body_size 12m;
50
+ proxy_read_timeout 1800s;
51
+ proxy_send_timeout 1800s;
52
+ }
53
+ ```
54
+
55
+ Apply an edge rate limit appropriate to available RAM and latency. One model worker is deliberate: a concurrent long-context generation can overcommit local memory. Scale by adding isolated model replicas behind a router, not by increasing Uvicorn workers inside one model container.
56
+
57
+ ## Authentication and CORS
58
+
59
+ Set `PHILLNET_API_KEY` in `.env`. The protected completion route accepts either of these headers:
60
+
61
+ ```http
62
+ Authorization: Bearer YOUR_SECRET
63
+ ```
64
+
65
+ ```http
66
+ X-API-Key: YOUR_SECRET
67
+ ```
68
+
69
+ Set `PHILLNET_CORS_ORIGINS` to a comma-separated list of known origins such as `https://app.example.com`. Leave it empty for server-to-server callers. Never use a wildcard origin together with a browser-facing credential policy.
70
+
71
+ ## Inference and input limits
72
+
73
+ | Guardrail | Default | Reason |
74
+ |---|---:|---|
75
+ | Concurrent generations | 1 | Avoids contention and memory overcommit on one local checkpoint. |
76
+ | Maximum visible answer | 8,192 tokens | Matches the model’s persisted response policy. |
77
+ | Maximum images per request | 4 | Bounds multimodal preprocessing work. |
78
+ | Maximum decoded image size | 10 MiB | Prevents oversized inline payloads. |
79
+ | Maximum image pixels | 24,000,000 | Mitigates decompression-bomb and preprocessing risks. |
80
+ | Maximum request body | 12 MiB | Rejects oversized JSON before inference. |
81
+
82
+ ## Health and observability
83
+
84
+ Use these endpoints in the surrounding platform:
85
+
86
+ | Endpoint | Purpose | Authentication |
87
+ |---|---|---|
88
+ | `GET /health` | Basic liveness and advertised capability surface. | No |
89
+ | `GET /ready` | Readiness after checkpoint and processor load. | No |
90
+ | `POST /v1/chat/completions` | Text and optional still-image inference. | Required when `PHILLNET_API_KEY` is set. |
91
+
92
+ The completion response includes `usage` counts and `elapsed_seconds`, which are sufficient for basic application-side request telemetry. Do not log request images or full user prompts unless your data-retention policy explicitly allows it.
93
+
94
+ ## Upgrade and rollback
95
+
96
+ Build tagged images rather than relying on mutable local source.
97
+
98
+ ```bash
99
+ docker compose build
100
+ docker image tag phillnet-mini-text-vision:1.1.0 registry.example.com/phillnet-mini-text-vision:1.1.0
101
+ # Push to your trusted registry, then deploy that immutable tag.
102
+ ```
103
+
104
+ Retain the preceding image tag and the matching `RELEASE_MANIFEST.json`. Verify the checkpoint hash before switching traffic. Roll back by returning the compose image reference to the prior known-good tag and running `docker compose up -d`.
105
+
106
+ ## Scope boundary
107
+
108
+ This service supports only **text generation** and **still-image understanding**. It does not expose image generation, video generation, audio, tools, agents, browsing, remote URL fetching, or arbitrary local file access.
README.md ADDED
@@ -0,0 +1,285 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ language:
4
+ - en
5
+ pipeline_tag: image-text-to-text
6
+ tags:
7
+ - text-generation
8
+ - visual-question-answering
9
+ - multimodal
10
+ - custom-code
11
+ - adaptive-reasoning
12
+ - text-vision-only
13
+ library_name: transformers
14
+ ---
15
+
16
+ # Phillnet Mini Text-Vision
17
+
18
+ **Phillnet Mini Text-Vision** is a lean, self-contained derivative of [ayjays132/Phillnet-Mini-Omni-Max][1] for two focused capabilities: **language generation** and **still-image understanding**. It preserves the source checkpoint’s language and visual-encoder path while removing the SDXL synthesis stack and nonessential runtime surfaces.
19
+
20
+ > **Designed for:** text completion, multimodal chat, visual question answering, image-grounded reasoning, and code or HTML generation.
21
+ > **Not included:** image generation, video generation, audio, agents, tools, SDXL, diffusion schedulers, U-Net, VAE, or SDXL text encoders.
22
+
23
+ ## Example gallery
24
+
25
+ | Direct example: Horizon Tasks | Medium example: Solara Grid | Maximum adaptive example: Orbit Notes |
26
+ |---|---|---|
27
+ | [![Horizon Tasks landing page](https://private-us-east-1.manuscdn.com/sessionFile/BqAFZMqxyrtm3s5wdTv4jM/sandbox/HtU7EIGVyL1k8mG5rg24G6-images_1787379131756_na1fn_L2hvbWUvdWJ1bnR1L3BoaWxsbmV0X3dvcmsvUGhpbGxuZXQtTWluaS1UZXh0LVZpc2lvbi9leGFtcGxlcy9kaXJlY3QtaG9yaXpvbi10YXNrcw.png?Policy=eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9wcml2YXRlLXVzLWVhc3QtMS5tYW51c2Nkbi5jb20vc2Vzc2lvbkZpbGUvQnFBRlpNcXh5cnRtM3M1d2RUdjRqTS9zYW5kYm94L0h0VTdFSUdWeUwxazhtRzVyZzI0RzYtaW1hZ2VzXzE3ODczNzkxMzE3NTZfbmExZm5fTDJodmJXVXZkV0oxYm5SMUwzQm9hV3hzYm1WMFgzZHZjbXN2VUdocGJHeHVaWFF0VFdsdWFTMVVaWGgwTFZacGMybHZiaTlsZUdGdGNHeGxjeTlrYVhKbFkzUXRhRzl5YVhwdmJpMTBZWE5yY3cucG5nIiwiQ29uZGl0aW9uIjp7IkRhdGVMZXNzVGhhbiI6eyJBV1M6RXBvY2hUaW1lIjoxNzg5NDMwNDAwfX19XX0_&Key-Pair-Id=K2QY5QTL8JSY6C&Signature=MEYCIQD97wMIU0b2gCDDgui8twaZctOYKGp9qOHQvycgpDIv7AIhAN01ndIde03tb6BnOGz4sq0nDhhTXpu5IOnk-r9mrTxj)](examples/direct-horizon-tasks.html) | [![Solara Grid dashboard](https://private-us-east-1.manuscdn.com/sessionFile/BqAFZMqxyrtm3s5wdTv4jM/sandbox/HtU7EIGVyL1k8mG5rg24G6-images_1787379131756_na1fn_L2hvbWUvdWJ1bnR1L3BoaWxsbmV0X3dvcmsvUGhpbGxuZXQtTWluaS1UZXh0LVZpc2lvbi9leGFtcGxlcy9tZWRpdW0tc29sYXJhLWRhc2hib2FyZA.png?Policy=eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9wcml2YXRlLXVzLWVhc3QtMS5tYW51c2Nkbi5jb20vc2Vzc2lvbkZpbGUvQnFBRlpNcXh5cnRtM3M1d2RUdjRqTS9zYW5kYm94L0h0VTdFSUdWeUwxazhtRzVyZzI0RzYtaW1hZ2VzXzE3ODczNzkxMzE3NTZfbmExZm5fTDJodmJXVXZkV0oxYm5SMUwzQm9hV3hzYm1WMFgzZHZjbXN2VUdocGJHeHVaWFF0VFdsdWFTMVVaWGgwTFZacGMybHZiaTlsZUdGdGNHeGxjeTl0WldScGRXMHRjMjlzWVhKaExXUmhjMmhpYjJGeVpBLnBuZyIsIkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTc4OTQzMDQwMH19fV19&Key-Pair-Id=K2QY5QTL8JSY6C&Signature=MEYCIQCIqMN-TCflATcFV0J0Zto7HWLe~wz0OyJgsWnEgqXCLAIhAIUdO7UE5MbCIw96dBisgVbq6~SMMwtlgfSPYdxXpxMt)](examples/medium-solara-dashboard.html) | [![Orbit Notes landing page](https://private-us-east-1.manuscdn.com/sessionFile/BqAFZMqxyrtm3s5wdTv4jM/sandbox/HtU7EIGVyL1k8mG5rg24G6-images_1787379131756_na1fn_L2hvbWUvdWJ1bnR1L3BoaWxsbmV0X3dvcmsvUGhpbGxuZXQtTWluaS1UZXh0LVZpc2lvbi9leGFtcGxlcy9tYXgtb3JiaXQtbm90ZXM.png?Policy=eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9wcml2YXRlLXVzLWVhc3QtMS5tYW51c2Nkbi5jb20vc2Vzc2lvbkZpbGUvQnFBRlpNcXh5cnRtM3M1d2RUdjRqTS9zYW5kYm94L0h0VTdFSUdWeUwxazhtRzVyZzI0RzYtaW1hZ2VzXzE3ODczNzkxMzE3NTZfbmExZm5fTDJodmJXVXZkV0oxYm5SMUwzQm9hV3hzYm1WMFgzZHZjbXN2VUdocGJHeHVaWFF0VFdsdWFTMVVaWGgwTFZacGMybHZiaTlsZUdGdGNHeGxjeTl0WVhndGIzSmlhWFF0Ym05MFpYTS5wbmciLCJDb25kaXRpb24iOnsiRGF0ZUxlc3NUaGFuIjp7IkFXUzpFcG9jaFRpbWUiOjE3ODk0MzA0MDB9fX1dfQ__&Key-Pair-Id=K2QY5QTL8JSY6C&Signature=MEQCIGoHLMxYVWPWicHZXdTNzYdrNk9LW4wr0vMsNzWp5yGcAiAbPbTwfQPNpkHRNTFvNLCC7bZRbjXRyG8uvjrr2q-uVQ__)](examples/max-orbit-notes.html) |
28
+ | [Open HTML source](examples/direct-horizon-tasks.html) | [Open HTML source](examples/medium-solara-dashboard.html) | [Open HTML source](examples/max-orbit-notes.html) |
29
+
30
+ The example pages are standalone HTML files with inline CSS and JavaScript. They are included as reproducible reference outputs for landing-page, dashboard, and long-form one-shot generation workflows.
31
+
32
+ ## What changed from the omni source
33
+
34
+ The original source model includes an SDXL-backed synthesis route. This package excludes that route at the file, configuration, and runtime levels so the deployed capability surface remains intentionally narrow.
35
+
36
+ | Area | Phillnet Mini Text-Vision behavior |
37
+ |---|---|
38
+ | Language generation | **Enabled** through `DendroForCausalLM.generate(...)`. |
39
+ | Still-image understanding | **Enabled** through the local `DendroVisionProcessor`, the visual encoder, and multimodal chat inputs. |
40
+ | SDXL image synthesis | **Removed**, including SDXL weights, U-Net, VAE, text encoders, scheduler, adapter, and public generation methods. |
41
+ | Video synthesis | **Removed**. |
42
+ | Audio, tools, agents, orchestration | **Not exposed** by the lean package or the included service. |
43
+ | External processor requirement | **Removed**. The package includes a local processor that uses the checkpoint’s native Qwen tokenizer and vision-token contract. |
44
+
45
+ The retained `model.safetensors` checkpoint is unchanged from the source core checkpoint:
46
+
47
+ ```text
48
+ SHA-256: f1a913f99f8ce921c1aa79982a09eaee6744448766f09a4756751e2d4b9342fc
49
+ ```
50
+
51
+ ## Adaptive reasoning defaults
52
+
53
+ The package is configured for deep adaptive inference by default. The model’s private reasoning stays private; only the final answer is returned from `model.generate(...)`.
54
+
55
+ | Default | Value | Meaning |
56
+ |---|---:|---|
57
+ | Default reasoning effort | `max` | Uses the highest configured adaptive effort profile when the caller does not override it. |
58
+ | Adaptive private reasoning | `true` | The private phase can expand and stop naturally at the model’s end-of-thinking token. |
59
+ | Private-reasoning policy | `adaptive_context` | The private phase is bounded by the active physical context/cache boundary rather than a small caller cap. |
60
+ | Visible answer ceiling | **8,192 tokens** | Keeps long-form answers and code generation practical and predictable. |
61
+ | Active context/cache window | 32,768 tokens | Establishes the physical boundary for prompt, private reasoning, and answer allocation. |
62
+
63
+ > “Adaptive” does not mean unbounded hardware usage. The model may reason for as long as the active context and cache allocation permit, then reserves a separately bounded visible answer. This package’s visible-answer ceiling is intentionally fixed at **8,192 tokens**.
64
+
65
+ The included FastAPI service also defaults to `reasoning_effort="max"` and `max_tokens=8192`. Requests can explicitly select `direct`, `low`, `medium`, `high`, or `max` when latency or cost behavior needs to differ.
66
+
67
+ ## Installation
68
+
69
+ ```bash
70
+ pip install -r requirements.txt
71
+ ```
72
+
73
+ The package uses custom Transformers code, so load it from a local checkout or trusted repository with `trust_remote_code=True` for the model and processor.
74
+
75
+ ```python
76
+ import torch
77
+ from transformers import AutoModelForCausalLM, AutoProcessor
78
+
79
+ model_dir = "./Phillnet-Mini-Text-Vision"
80
+
81
+ processor = AutoProcessor.from_pretrained(
82
+ model_dir,
83
+ trust_remote_code=True,
84
+ )
85
+ model = AutoModelForCausalLM.from_pretrained(
86
+ model_dir,
87
+ trust_remote_code=True,
88
+ dtype=torch.bfloat16,
89
+ low_cpu_mem_usage=True,
90
+ ).eval()
91
+ ```
92
+
93
+ ## Text generation
94
+
95
+ ### Use the persisted maximum adaptive defaults
96
+
97
+ The call below deliberately omits `max_new_tokens`, `max_reasoning_tokens`, and `reasoning_budget`. The model therefore uses its loaded maximum adaptive profile and 8,192-token visible-answer ceiling.
98
+
99
+ ```python
100
+ messages = [{
101
+ "role": "user",
102
+ "content": "Write a concise release note for a focused note-taking application.",
103
+ }]
104
+
105
+ inputs = processor.apply_chat_template(
106
+ messages,
107
+ tokenize=True,
108
+ add_generation_prompt=True,
109
+ return_dict=True,
110
+ return_tensors="pt",
111
+ )
112
+
113
+ output = model.generate(
114
+ **inputs,
115
+ do_sample=False,
116
+ use_cache=True,
117
+ )
118
+
119
+ prompt_length = inputs["input_ids"].shape[1]
120
+ answer = processor.tokenizer.decode(
121
+ output[0, prompt_length:],
122
+ skip_special_tokens=True,
123
+ )
124
+ print(answer)
125
+ ```
126
+
127
+ ### Choose a lower-latency profile
128
+
129
+ ```python
130
+ output = model.generate(
131
+ **inputs,
132
+ reasoning_effort="direct",
133
+ max_new_tokens=256,
134
+ do_sample=False,
135
+ use_cache=True,
136
+ )
137
+ ```
138
+
139
+ ## Still-image understanding
140
+
141
+ The bundled `DendroVisionProcessor` accepts PIL images directly. It performs local visual patch preparation, emits the correct visual placeholder sequence, and injects visual features before the assistant response prefix.
142
+
143
+ ```python
144
+ from PIL import Image
145
+
146
+ image = Image.open("scene.png").convert("RGB")
147
+ messages = [{
148
+ "role": "user",
149
+ "content": [
150
+ {"type": "image", "image": image},
151
+ {"type": "text", "text": "Describe the important objects and their relative positions."},
152
+ ],
153
+ }]
154
+
155
+ inputs = processor.apply_chat_template(
156
+ messages,
157
+ tokenize=True,
158
+ add_generation_prompt=True,
159
+ return_dict=True,
160
+ return_tensors="pt",
161
+ )
162
+
163
+ output = model.generate(
164
+ **inputs,
165
+ reasoning_effort="medium",
166
+ max_new_tokens=512,
167
+ do_sample=False,
168
+ use_cache=True,
169
+ )
170
+
171
+ prompt_length = inputs["input_ids"].shape[1]
172
+ print(processor.tokenizer.decode(output[0, prompt_length:], skip_special_tokens=True))
173
+ ```
174
+
175
+ ## HTTP service
176
+
177
+ The package includes a small OpenAI-style FastAPI service that intentionally exposes only text generation and still-image understanding.
178
+
179
+ ```bash
180
+ pip install -r requirements.txt -r requirements-server.txt
181
+ uvicorn server:app --host 0.0.0.0 --port 8000
182
+ ```
183
+
184
+ Check availability:
185
+
186
+ ```bash
187
+ curl http://localhost:8000/health
188
+ ```
189
+
190
+ Submit a text request using the default maximum adaptive behavior:
191
+
192
+ ```bash
193
+ curl http://localhost:8000/v1/chat/completions \
194
+ -H 'Content-Type: application/json' \
195
+ --data '{
196
+ "messages": [{"role": "user", "content": "Create a compact project brief for a solar dashboard."}]
197
+ }'
198
+ ```
199
+
200
+ The endpoint accepts a base64-encoded still image through an `image_base64` content item. It deliberately does not fetch remote image URLs.
201
+
202
+ ```bash
203
+ IMAGE_B64=$(base64 -w0 scene.png)
204
+ curl http://localhost:8000/v1/chat/completions \
205
+ -H 'Content-Type: application/json' \
206
+ --data "{
207
+ \"messages\": [{
208
+ \"role\": \"user\",
209
+ \"content\": [
210
+ {\"type\": \"image\", \"image_base64\": \"${IMAGE_B64}\"},
211
+ {\"type\": \"text\", \"text\": \"What is visible in this image?\"}
212
+ ]
213
+ }]
214
+ }"
215
+ ```
216
+
217
+ ## Production deployment
218
+
219
+ The repository includes a hardened single-worker container path for self-hosted inference. It runs as a non-root user, exposes readiness checks, serializes local generation to prevent memory overcommit, accepts only inline base64 still images, and supports optional bearer or `X-API-Key` authentication.
220
+
221
+ ```bash
222
+ # On a server that has Docker and Docker Compose installed.
223
+ cp .env.example .env
224
+ # Edit .env: set a long PHILLNET_API_KEY and your exact browser origin.
225
+ docker compose up --build -d
226
+
227
+ docker compose ps
228
+ curl http://127.0.0.1:8000/health
229
+ ```
230
+
231
+ The compose file binds to `127.0.0.1:8000` by default. Place a TLS-terminating reverse proxy in front of it for public traffic, set `PHILLNET_CORS_ORIGINS` to the exact allowed application origin, and keep `PHILLNET_API_KEY` non-empty. Do not place the model API directly on the public internet without authentication and rate limiting at the edge.
232
+
233
+ | Environment variable | Default | Deployment purpose |
234
+ |---|---|---|
235
+ | `PHILLNET_API_KEY` | Empty locally | Enforces bearer or `X-API-Key` authentication when set. |
236
+ | `PHILLNET_CORS_ORIGINS` | Disabled | Allows only named browser origins. |
237
+ | `PHILLNET_MAX_IMAGE_BYTES` | 10 MiB | Rejects oversized decoded image payloads. |
238
+ | `PHILLNET_MAX_IMAGE_PIXELS` | 24,000,000 | Limits decoded image dimensions. |
239
+ | `PHILLNET_MAX_REQUEST_BYTES` | 12 MiB | Rejects oversized request bodies before inference. |
240
+ | `PHILLNET_BIND_ADDRESS` | `127.0.0.1` | Keeps the compose service loopback-bound by default. |
241
+ | `PHILLNET_PORT` | `8000` | Changes the host listener port. |
242
+
243
+ See [HF_UPLOAD.md](HF_UPLOAD.md) for the complete Hugging Face upload workflow and [`.env.example`](.env.example) for deployment values.
244
+
245
+ ## Validation snapshot
246
+
247
+ The lean package was tested with structural checks, local model loading, text completion, visual patch preparation, visual grounding, basic spatial grounding, and service-level requests.
248
+
249
+ | Validation | Verified result |
250
+ |---|---|
251
+ | Core checkpoint loading | `DendroForCausalLM` loaded in BF16. |
252
+ | Tokenizer contract | Local `Qwen2Tokenizer` matches the checkpoint’s 248,320-token vocabulary. |
253
+ | Processor contract | `AutoProcessor` resolves to `DendroVisionProcessor`. |
254
+ | Text coherence | Deterministic chat probes returned `Paris`, `4`, and `blue`. |
255
+ | Color grounding | Solid red, green, blue, white, and black probes returned the corresponding color. |
256
+ | Spatial grounding | A red-left / blue-right scene returned `red` for left and `blue` for right. |
257
+ | SDXL exclusion | SDXL weights, code, configuration, dependencies, and synthesis APIs are absent. |
258
+ | Full adaptive HTML generation | A complete 2,991-token Orbit Notes page was generated after 2,944 private-reasoning tokens. |
259
+
260
+ ## Layout
261
+
262
+ | Path | Purpose |
263
+ |---|---|
264
+ | `model.safetensors` | Retained language and visual-understanding checkpoint. |
265
+ | `processing_dendro_omni.py` | Local native-tokenizer text-and-image processor. |
266
+ | `modeling_dendro_omni.py` | Custom language and multimodal generation runtime. |
267
+ | `server.py` | Text-and-vision-only HTTP service. |
268
+ | `examples/` | Completed HTML examples and PNG previews. |
269
+ | `ADAPTIVE_DEFAULTS_REPORT.md` | Deep configuration and long-form generation report. |
270
+ | `PRODUCTION.md` | Deployment architecture, authentication, limits, reverse proxy, and operating guide. |
271
+ | `HF_UPLOAD.md` | Current Hugging Face upload workflow and verification checklist. |
272
+ | `COHERENCE_REPAIR_REPORT.md` | Tokenizer, chat-template, and spatial-input repair record. |
273
+ | `IMPLEMENTATION_REPORT.md` | Original lean-package removal and validation report. |
274
+
275
+ ## Operational notes
276
+
277
+ The core checkpoint is approximately 1.76 GB. Deep adaptive reasoning and long code generation can require substantially more memory and time than direct short-answer requests. For production use, evaluate target hardware, concurrency, request limits, observability, and authentication before exposing the API publicly.
278
+
279
+ ## License and attribution
280
+
281
+ This derivative retains the upstream model’s Apache-2.0 license designation. Consult the source model card and included licensing files before redistribution or use in a product. [1]
282
+
283
+ ## References
284
+
285
+ [1]: https://huggingface.co/ayjays132/Phillnet-Mini-Omni-Max "ayjays132/Phillnet-Mini-Omni-Max"
RELEASE_MANIFEST.json ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "schema_version": "1.0",
3
+ "release": {
4
+ "name": "Phillnet Mini Text-Vision",
5
+ "version": "1.1.0",
6
+ "release_date": "2026-08-21",
7
+ "license": "Apache-2.0",
8
+ "source_model": "https://huggingface.co/ayjays132/Phillnet-Mini-Omni-Max"
9
+ },
10
+ "repository": {
11
+ "upload_root": true,
12
+ "file_count_excluding_this_manifest": 48,
13
+ "size_bytes_excluding_this_manifest": 1785015195,
14
+ "lfs_artifacts": [
15
+ "model.safetensors",
16
+ "tokenizer.json"
17
+ ]
18
+ },
19
+ "artifacts": {
20
+ "model.safetensors": {
21
+ "bytes": 1763655304,
22
+ "sha256": "f1a913f99f8ce921c1aa79982a09eaee6744448766f09a4756751e2d4b9342fc"
23
+ },
24
+ "tokenizer.json": {
25
+ "bytes": 19989325,
26
+ "sha256": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523"
27
+ }
28
+ },
29
+ "capabilities": {
30
+ "enabled": [
31
+ "text-generation",
32
+ "still-image-understanding"
33
+ ],
34
+ "disabled": [
35
+ "sdxl-image-generation",
36
+ "video-generation",
37
+ "audio",
38
+ "tools",
39
+ "agents",
40
+ "remote-image-url-fetching"
41
+ ]
42
+ },
43
+ "inference_defaults": {
44
+ "reasoning_effort": "max",
45
+ "adaptive_private_reasoning": true,
46
+ "private_reasoning_policy": "adaptive_context",
47
+ "max_visible_answer_tokens": 8192,
48
+ "active_context_window_tokens": 32768
49
+ },
50
+ "deployment": {
51
+ "container": "Dockerfile",
52
+ "compose": "docker-compose.yml",
53
+ "service": "server.py",
54
+ "authentication": "PHILLNET_API_KEY (optional locally, required for public deployment)",
55
+ "api_route": "/v1/chat/completions",
56
+ "health_route": "/health",
57
+ "readiness_route": "/ready"
58
+ }
59
+ }
RELEASE_VALIDATION.md ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Release Validation
2
+
3
+ **Release:** Phillnet Mini Text-Vision v1.1.0
4
+ **Date:** 21 August 2026
5
+
6
+ ## Completed checks
7
+
8
+ | Check | Result |
9
+ |---|---|
10
+ | Custom model Python modules | All root Python modules compile successfully. |
11
+ | `config.json` and `RELEASE_MANIFEST.json` | Valid JSON. |
12
+ | Retained model checkpoint | SHA-256 matches the value in `RELEASE_MANIFEST.json`. |
13
+ | Example gallery | All three HTML and PNG assets are present. |
14
+ | Hugging Face large-file policy | `model.safetensors` and `tokenizer.json` are explicitly marked for LFS. |
15
+ | Repository hygiene | `.gitignore` excludes bytecode caches, secrets, logs, and archives. |
16
+ | Production files | `Dockerfile`, `docker-compose.yml`, `.env.example`, `PRODUCTION.md`, and `HF_UPLOAD.md` are present. |
17
+ | API-key gate | Missing credentials are rejected when `PHILLNET_API_KEY` is set; bearer and `X-API-Key` forms are accepted. |
18
+ | Request safeguards | Maximum image byte, pixel, and request-body limits were validated. |
19
+ | Temporary deployment | Fresh `/health` and `/ready` checks returned a ready text-and-vision-only service. |
20
+
21
+ ## Validation boundary
22
+
23
+ The available environment does not include Docker, so the container image and Compose runtime were inspected and validated statically rather than built locally. The Python service was compiled and loaded directly with the full model checkpoint, then checked through its HTTP health and readiness routes.
24
+
25
+ ## Upload status
26
+
27
+ The repository is ready to upload as a complete Hugging Face model repository. Follow [HF_UPLOAD.md](HF_UPLOAD.md) from the parent directory and verify the checksum in [RELEASE_MANIFEST.json](RELEASE_MANIFEST.json) after transfer.
__init__.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Phillnet Mini Text-Vision: language generation and image understanding only."""
2
+
3
+ from .cache import DendroKVCache, QuantizedCacheTensor
4
+ from .configuration_dendro_omni import DendroOmniConfig
5
+ from .modeling_dendro_omni import (
6
+ DendroAnswerGenerationOutput,
7
+ DendroBaseModelOutputWithPast,
8
+ DendroCausalLMOutput,
9
+ DendroForCausalLM,
10
+ DendroOmniModel,
11
+ register_dendro_auto_classes,
12
+ )
13
+ from .processing_dendro_omni import DendroVisionProcessor
14
+ from .tokenization_dendro_omni import DendroByteTokenizer
15
+ from .vision_routing import DendroVisionAnswer, build_vision_views
16
+
17
+ __version__ = "1.0.0-text-vision"
18
+
19
+ __all__ = [
20
+ "DendroAnswerGenerationOutput",
21
+ "DendroBaseModelOutputWithPast",
22
+ "DendroByteTokenizer",
23
+ "DendroCausalLMOutput",
24
+ "DendroForCausalLM",
25
+ "DendroKVCache",
26
+ "DendroOmniConfig",
27
+ "DendroOmniModel",
28
+ "DendroVisionAnswer",
29
+ "DendroVisionProcessor",
30
+ "QuantizedCacheTensor",
31
+ "build_vision_views",
32
+ "register_dendro_auto_classes",
33
+ ]
chat_template.jinja ADDED
@@ -0,0 +1,154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- set image_count = namespace(value=0) %}
2
+ {%- set video_count = namespace(value=0) %}
3
+ {%- macro render_content(content, do_vision_count, is_system_content=false) %}
4
+ {%- if content is string %}
5
+ {{- content }}
6
+ {%- elif content is iterable and content is not mapping %}
7
+ {%- for item in content %}
8
+ {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
9
+ {%- if is_system_content %}
10
+ {{- raise_exception('System message cannot contain images.') }}
11
+ {%- endif %}
12
+ {%- if do_vision_count %}
13
+ {%- set image_count.value = image_count.value + 1 %}
14
+ {%- endif %}
15
+ {%- if add_vision_id %}
16
+ {{- 'Picture ' ~ image_count.value ~ ': ' }}
17
+ {%- endif %}
18
+ {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
19
+ {%- elif 'video' in item or item.type == 'video' %}
20
+ {%- if is_system_content %}
21
+ {{- raise_exception('System message cannot contain videos.') }}
22
+ {%- endif %}
23
+ {%- if do_vision_count %}
24
+ {%- set video_count.value = video_count.value + 1 %}
25
+ {%- endif %}
26
+ {%- if add_vision_id %}
27
+ {{- 'Video ' ~ video_count.value ~ ': ' }}
28
+ {%- endif %}
29
+ {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
30
+ {%- elif 'text' in item %}
31
+ {{- item.text }}
32
+ {%- else %}
33
+ {{- raise_exception('Unexpected item type in content.') }}
34
+ {%- endif %}
35
+ {%- endfor %}
36
+ {%- elif content is none or content is undefined %}
37
+ {{- '' }}
38
+ {%- else %}
39
+ {{- raise_exception('Unexpected content type.') }}
40
+ {%- endif %}
41
+ {%- endmacro %}
42
+ {%- if not messages %}
43
+ {{- raise_exception('No messages provided.') }}
44
+ {%- endif %}
45
+ {%- if tools and tools is iterable and tools is not mapping %}
46
+ {{- '<|im_start|>system\n' }}
47
+ {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
48
+ {%- for tool in tools %}
49
+ {{- "\n" }}
50
+ {{- tool | tojson }}
51
+ {%- endfor %}
52
+ {{- "\n</tools>" }}
53
+ {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
54
+ {%- if messages[0].role == 'system' %}
55
+ {%- set content = render_content(messages[0].content, false, true)|trim %}
56
+ {%- if content %}
57
+ {{- '\n\n' + content }}
58
+ {%- endif %}
59
+ {%- endif %}
60
+ {{- '<|im_end|>\n' }}
61
+ {%- else %}
62
+ {%- if messages[0].role == 'system' %}
63
+ {%- set content = render_content(messages[0].content, false, true)|trim %}
64
+ {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
65
+ {%- endif %}
66
+ {%- endif %}
67
+ {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
68
+ {%- for message in messages[::-1] %}
69
+ {%- set index = (messages|length - 1) - loop.index0 %}
70
+ {%- if ns.multi_step_tool and message.role == "user" %}
71
+ {%- set content = render_content(message.content, false)|trim %}
72
+ {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
73
+ {%- set ns.multi_step_tool = false %}
74
+ {%- set ns.last_query_index = index %}
75
+ {%- endif %}
76
+ {%- endif %}
77
+ {%- endfor %}
78
+ {%- if ns.multi_step_tool %}
79
+ {{- raise_exception('No user query found in messages.') }}
80
+ {%- endif %}
81
+ {%- for message in messages %}
82
+ {%- set content = render_content(message.content, true)|trim %}
83
+ {%- if message.role == "system" %}
84
+ {%- if not loop.first %}
85
+ {{- raise_exception('System message must be at the beginning.') }}
86
+ {%- endif %}
87
+ {%- elif message.role == "user" %}
88
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
89
+ {%- elif message.role == "assistant" %}
90
+ {%- set reasoning_content = '' %}
91
+ {%- if message.reasoning_content is string %}
92
+ {%- set reasoning_content = message.reasoning_content %}
93
+ {%- else %}
94
+ {%- if '</think>' in content %}
95
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
96
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
97
+ {%- endif %}
98
+ {%- endif %}
99
+ {%- set reasoning_content = reasoning_content|trim %}
100
+ {%- if loop.index0 > ns.last_query_index %}
101
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
102
+ {%- else %}
103
+ {{- '<|im_start|>' + message.role + '\n' + content }}
104
+ {%- endif %}
105
+ {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
106
+ {%- for tool_call in message.tool_calls %}
107
+ {%- if tool_call.function is defined %}
108
+ {%- set tool_call = tool_call.function %}
109
+ {%- endif %}
110
+ {%- if loop.first %}
111
+ {%- if content|trim %}
112
+ {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
113
+ {%- else %}
114
+ {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
115
+ {%- endif %}
116
+ {%- else %}
117
+ {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
118
+ {%- endif %}
119
+ {%- if tool_call.arguments is defined %}
120
+ {%- for args_name, args_value in tool_call.arguments|items %}
121
+ {{- '<parameter=' + args_name + '>\n' }}
122
+ {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
123
+ {{- args_value }}
124
+ {{- '\n</parameter>\n' }}
125
+ {%- endfor %}
126
+ {%- endif %}
127
+ {{- '</function>\n</tool_call>' }}
128
+ {%- endfor %}
129
+ {%- endif %}
130
+ {{- '<|im_end|>\n' }}
131
+ {%- elif message.role == "tool" %}
132
+ {%- if loop.previtem and loop.previtem.role != "tool" %}
133
+ {{- '<|im_start|>user' }}
134
+ {%- endif %}
135
+ {{- '\n<tool_response>\n' }}
136
+ {{- content }}
137
+ {{- '\n</tool_response>' }}
138
+ {%- if not loop.last and loop.nextitem.role != "tool" %}
139
+ {{- '<|im_end|>\n' }}
140
+ {%- elif loop.last %}
141
+ {{- '<|im_end|>\n' }}
142
+ {%- endif %}
143
+ {%- else %}
144
+ {{- raise_exception('Unexpected message role.') }}
145
+ {%- endif %}
146
+ {%- endfor %}
147
+ {%- if add_generation_prompt %}
148
+ {{- '<|im_start|>assistant\n' }}
149
+ {%- if enable_thinking is defined and enable_thinking is true %}
150
+ {{- '<think>\n' }}
151
+ {%- else %}
152
+ {{- '<think>\n\n</think>\n\n' }}
153
+ {%- endif %}
154
+ {%- endif %}
config.json ADDED
The diff for this file is too large to render. See raw diff
 
docker-compose.yml ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ services:
2
+ phillnet-mini-text-vision:
3
+ build:
4
+ context: .
5
+ dockerfile: Dockerfile
6
+ image: phillnet-mini-text-vision:1.1.0
7
+ restart: unless-stopped
8
+ env_file:
9
+ - .env
10
+ ports:
11
+ - "${PHILLNET_BIND_ADDRESS:-127.0.0.1}:${PHILLNET_PORT:-8000}:8000"
12
+ init: true
13
+ read_only: true
14
+ tmpfs:
15
+ - /tmp:rw,noexec,nosuid,size=256m
16
+ security_opt:
17
+ - no-new-privileges:true
18
+ healthcheck:
19
+ test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/ready', timeout=3)"]
20
+ interval: 30s
21
+ timeout: 5s
22
+ retries: 3
23
+ start_period: 120s
examples/direct-horizon-tasks.png ADDED

Git LFS Details

  • SHA256: b7e02ffd5c35aa87da5f7e3d0dafeeb0db79123deb24da94c43b1d761df7df11
  • Pointer size: 131 Bytes
  • Size of remote file: 103 kB
examples/max-orbit-notes.png ADDED
examples/medium-solara-dashboard.png ADDED

Git LFS Details

  • SHA256: 98a236ecd574fe28ec51f76b4525602a2964551f53f7d560554c90475629401b
  • Pointer size: 131 Bytes
  • Size of remote file: 173 kB
modalities.py ADDED
@@ -0,0 +1,407 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Real multimodal patchification and packing into one Dendro token space."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass
6
+ from typing import Any
7
+
8
+ import torch
9
+ from torch.nn import functional as F
10
+
11
+ from ._source_bound import SourceBoundModule
12
+ from .configuration_dendro_omni import DendroOmniConfig
13
+ from .source import DendroSourceLayer
14
+ from .spatial import (
15
+ MODALITY_AUDIO,
16
+ MODALITY_IMAGE,
17
+ MODALITY_POSITION_OFFSETS,
18
+ MODALITY_SENSOR,
19
+ MODALITY_TEXT,
20
+ MODALITY_VIDEO,
21
+ DendroSpatialEncoder,
22
+ )
23
+
24
+
25
+ @dataclass(slots=True)
26
+ class ModalitySegment:
27
+ name: str
28
+ modality_id: int
29
+ start: int
30
+ end: int
31
+ shape: tuple[int, ...]
32
+
33
+ @property
34
+ def length(self) -> int:
35
+ return self.end - self.start
36
+
37
+
38
+ @dataclass(slots=True)
39
+ class DendroModalityLayout:
40
+ modality_ids: torch.Tensor
41
+ sequence_positions: torch.Tensor
42
+ logical_positions: torch.Tensor
43
+ coordinates: torch.Tensor
44
+ is_prefix: torch.Tensor
45
+ attention_mask: torch.Tensor
46
+ segments: tuple[ModalitySegment, ...]
47
+ text_start: int
48
+ text_length: int
49
+
50
+ def to_dict(self) -> dict[str, Any]:
51
+ return {
52
+ "segments": [
53
+ {
54
+ "name": segment.name,
55
+ "modality_id": segment.modality_id,
56
+ "start": segment.start,
57
+ "end": segment.end,
58
+ "length": segment.length,
59
+ "shape": segment.shape,
60
+ }
61
+ for segment in self.segments
62
+ ],
63
+ "text_start": self.text_start,
64
+ "text_length": self.text_length,
65
+ "total_length": int(self.modality_ids.shape[-1]),
66
+ "prefix_length": int(self.is_prefix[0].sum().item()) if self.is_prefix.numel() else 0,
67
+ }
68
+
69
+
70
+ @dataclass(slots=True)
71
+ class DendroPackedInput:
72
+ hidden_states: torch.Tensor
73
+ layout: DendroModalityLayout
74
+ aligned_labels: torch.Tensor | None = None
75
+
76
+
77
+ class DendroOmniInputProjector(SourceBoundModule):
78
+ """Parameterless modality adapters backed entirely by ``DendroSourceLayer``."""
79
+
80
+ def __init__(self, config: DendroOmniConfig, source: DendroSourceLayer) -> None:
81
+ super().__init__(source)
82
+ self.config = config
83
+ self.spatial_encoder = DendroSpatialEncoder(config, source)
84
+
85
+ @staticmethod
86
+ def _batch_size(*values: torch.Tensor | None) -> int:
87
+ batches = [int(value.shape[0]) for value in values if value is not None]
88
+ if not batches:
89
+ raise ValueError("At least one text, image, audio, video, or sensor input is required")
90
+ if any(batch != batches[0] for batch in batches):
91
+ raise ValueError(f"All modalities must share a batch size, got {batches}")
92
+ return batches[0]
93
+
94
+ @staticmethod
95
+ def _normalize_coords(index: torch.Tensor, maximum: int) -> torch.Tensor:
96
+ if maximum <= 1:
97
+ return torch.zeros_like(index, dtype=torch.float32)
98
+ return index.float() / float(maximum - 1) * 2.0 - 1.0
99
+
100
+ def _token_features(self, input_ids: torch.Tensor, token_hidden: torch.Tensor) -> torch.Tensor:
101
+ source = self.source
102
+ offset = self.config.byte_offset
103
+ byte = input_ids - offset
104
+ atom_ids = torch.zeros_like(input_ids)
105
+ atom_ids = torch.where((byte >= ord("0")) & (byte <= ord("9")), 1, atom_ids)
106
+ atom_ids = torch.where(
107
+ ((byte >= ord("A")) & (byte <= ord("Z"))) | ((byte >= ord("a")) & (byte <= ord("z"))),
108
+ 2,
109
+ atom_ids,
110
+ )
111
+ atom_ids = torch.where((byte == 9) | (byte == 10) | (byte == 13) | (byte == 32), 3, atom_ids)
112
+ atom_ids = torch.where((byte >= 128) & (byte <= 255), 4, atom_ids)
113
+ atom_ids = torch.where(input_ids < offset, 5, atom_ids)
114
+ atoms = source.embedding(atom_ids, "token/atoms", 6, self.config.hidden_size)
115
+
116
+ # Atom -> bond -> molecule composition is deliberately pointwise here. Any
117
+ # cross-token neighborhood operation belongs inside cache-aware attention;
118
+ # otherwise a one-token decode chunk would not match full-sequence training.
119
+ bond_input = token_hidden * torch.tanh(atoms)
120
+ bonds = source.project(bond_input, "token/bonds", self.config.hidden_size, low_bit=False)
121
+ molecule_input = F.silu(bonds) + 0.5 * token_hidden + 0.25 * atoms
122
+ molecules = source.project(molecule_input, "token/molecules", self.config.hidden_size, low_bit=False)
123
+ gate = source.gate(torch.cat([token_hidden, atoms], dim=-1), "token/compose_gate", self.config.hidden_size)
124
+ return token_hidden + 0.15 * atoms + 0.10 * gate * bonds + 0.10 * (1.0 - gate) * molecules
125
+
126
+ def _text(
127
+ self,
128
+ input_ids: torch.Tensor | None,
129
+ inputs_embeds: torch.Tensor | None,
130
+ attention_mask: torch.Tensor | None,
131
+ *,
132
+ position_start: int,
133
+ prefix_mask: torch.Tensor | None,
134
+ ) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
135
+ if input_ids is None and inputs_embeds is None:
136
+ return None
137
+ if input_ids is not None and inputs_embeds is not None:
138
+ raise ValueError("Pass input_ids or inputs_embeds, not both")
139
+ if inputs_embeds is None:
140
+ assert input_ids is not None
141
+ hidden = self.source.embedding(
142
+ input_ids,
143
+ "token",
144
+ self.config.vocab_size,
145
+ self.config.hidden_size,
146
+ )
147
+ hidden = self._token_features(input_ids, hidden)
148
+ else:
149
+ if inputs_embeds.shape[-1] != self.config.hidden_size:
150
+ raise ValueError("inputs_embeds last dimension must equal hidden_size")
151
+ hidden = inputs_embeds
152
+ batch, length = hidden.shape[:2]
153
+ device = hidden.device
154
+ positions = torch.arange(position_start, position_start + length, device=device).expand(batch, -1)
155
+ logical = positions + MODALITY_POSITION_OFFSETS[MODALITY_TEXT]
156
+ coords = torch.zeros(batch, length, 4, device=device, dtype=hidden.dtype)
157
+ # Absolute coordinates must be invariant to chunking. Normalizing by the
158
+ # current input length made a token receive different spatial features during
159
+ # full-sequence training and cached one-token decoding.
160
+ denominator = max(1, self.config.max_position_embeddings - 1)
161
+ coords[..., 0] = (positions.to(hidden.dtype) / denominator * 2.0 - 1.0).clamp(-1.0, 1.0)
162
+ modality = torch.full((batch, length), MODALITY_TEXT, device=device, dtype=torch.long)
163
+ is_prefix = (
164
+ prefix_mask.to(device=device, dtype=torch.bool)
165
+ if prefix_mask is not None
166
+ else torch.zeros(batch, length, device=device, dtype=torch.bool)
167
+ )
168
+ mask = (
169
+ attention_mask.to(device=device, dtype=torch.bool)
170
+ if attention_mask is not None
171
+ else torch.ones(batch, length, device=device, dtype=torch.bool)
172
+ )
173
+ return hidden, {
174
+ "modality": modality,
175
+ "positions": positions,
176
+ "logical": logical,
177
+ "coords": coords,
178
+ "prefix": is_prefix,
179
+ "mask": mask,
180
+ }, (length,)
181
+
182
+ def _image(self, pixel_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
183
+ if pixel_values is None:
184
+ return None
185
+ if pixel_values.ndim != 4:
186
+ raise ValueError("pixel_values must be [batch, channels, height, width]")
187
+ batch, channels, height, width = pixel_values.shape
188
+ if channels != self.config.image_channels:
189
+ raise ValueError(f"Expected {self.config.image_channels} image channels, got {channels}")
190
+ patch = self.config.image_patch_size
191
+ pad_h, pad_w = (-height) % patch, (-width) % patch
192
+ values = F.pad(pixel_values, (0, pad_w, 0, pad_h))
193
+ grid_h, grid_w = values.shape[-2] // patch, values.shape[-1] // patch
194
+ patches = F.unfold(values, kernel_size=patch, stride=patch).transpose(1, 2)
195
+ hidden = self.source.project(patches, "modality/image_patch", self.config.hidden_size)
196
+ length = hidden.shape[1]
197
+ y = torch.arange(grid_h, device=hidden.device).repeat_interleave(grid_w)
198
+ x = torch.arange(grid_w, device=hidden.device).repeat(grid_h)
199
+ coords = torch.zeros(batch, length, 4, device=hidden.device, dtype=hidden.dtype)
200
+ coords[..., 1] = self._normalize_coords(y, grid_h)
201
+ coords[..., 2] = self._normalize_coords(x, grid_w)
202
+ local = torch.arange(length, device=hidden.device).expand(batch, -1)
203
+ return hidden, self._metadata(local, coords, MODALITY_IMAGE, prefix=True), (grid_h, grid_w)
204
+
205
+ def _audio(self, audio_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
206
+ if audio_values is None:
207
+ return None
208
+ if audio_values.ndim == 3:
209
+ audio_values = audio_values.mean(dim=1)
210
+ if audio_values.ndim != 2:
211
+ raise ValueError("audio_values must be [batch, samples] or [batch, channels, samples]")
212
+ patch, stride = self.config.audio_patch_size, self.config.audio_patch_stride
213
+ if audio_values.shape[-1] < patch:
214
+ audio_values = F.pad(audio_values, (0, patch - audio_values.shape[-1]))
215
+ remainder = (audio_values.shape[-1] - patch) % stride
216
+ if remainder:
217
+ audio_values = F.pad(audio_values, (0, stride - remainder))
218
+ windows = audio_values.unfold(-1, patch, stride)
219
+ hidden = self.source.project(windows, "modality/audio_patch", self.config.hidden_size)
220
+ length = hidden.shape[1]
221
+ local = torch.arange(length, device=hidden.device).expand(hidden.shape[0], -1)
222
+ coords = torch.zeros(hidden.shape[0], length, 4, device=hidden.device, dtype=hidden.dtype)
223
+ coords[..., 0] = self._normalize_coords(torch.arange(length, device=hidden.device), length)
224
+ # Frequency-energy coordinate gives raw wave patches a useful second axis.
225
+ spectrum = torch.fft.rfft(windows.float(), dim=-1).abs().mean(dim=-1)
226
+ spectrum = spectrum / spectrum.amax(dim=-1, keepdim=True).clamp_min(1e-8)
227
+ coords[..., 3] = spectrum.to(hidden.dtype) * 2.0 - 1.0
228
+ return hidden, self._metadata(local, coords, MODALITY_AUDIO, prefix=True), (length, patch)
229
+
230
+ def _video(self, video_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
231
+ if video_values is None:
232
+ return None
233
+ if video_values.ndim != 5:
234
+ raise ValueError("video_values must be [batch, frames, channels, height, width]")
235
+ batch, frames, channels, height, width = video_values.shape
236
+ if channels != self.config.image_channels:
237
+ raise ValueError(f"Expected {self.config.image_channels} video channels, got {channels}")
238
+ tube, patch = self.config.video_tubelet_size, self.config.video_patch_size
239
+ pad_t, pad_h, pad_w = (-frames) % tube, (-height) % patch, (-width) % patch
240
+ # F.pad follows reverse dimension order for [B,T,C,H,W].
241
+ values = F.pad(video_values, (0, pad_w, 0, pad_h, 0, 0, 0, pad_t))
242
+ tg, hg, wg = values.shape[1] // tube, values.shape[3] // patch, values.shape[4] // patch
243
+ blocks = values.reshape(batch, tg, tube, channels, hg, patch, wg, patch)
244
+ blocks = blocks.permute(0, 1, 4, 6, 2, 3, 5, 7).reshape(batch, tg * hg * wg, -1)
245
+ hidden = self.source.project(blocks, "modality/video_tubelet", self.config.hidden_size)
246
+ t = torch.arange(tg, device=hidden.device).repeat_interleave(hg * wg)
247
+ y = torch.arange(hg, device=hidden.device).repeat_interleave(wg).repeat(tg)
248
+ x = torch.arange(wg, device=hidden.device).repeat(hg * tg)
249
+ coords = torch.zeros(batch, hidden.shape[1], 4, device=hidden.device, dtype=hidden.dtype)
250
+ coords[..., 0] = self._normalize_coords(t, tg)
251
+ coords[..., 1] = self._normalize_coords(y, hg)
252
+ coords[..., 2] = self._normalize_coords(x, wg)
253
+ local = torch.arange(hidden.shape[1], device=hidden.device).expand(batch, -1)
254
+ return hidden, self._metadata(local, coords, MODALITY_VIDEO, prefix=True), (tg, hg, wg)
255
+
256
+ def _sensor(self, sensor_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
257
+ if sensor_values is None:
258
+ return None
259
+ if sensor_values.ndim == 2:
260
+ sensor_values = sensor_values.unsqueeze(1)
261
+ if sensor_values.ndim != 3:
262
+ raise ValueError("sensor_values must be [batch, steps, features] or [batch, features]")
263
+ target = self.config.sensor_feature_size
264
+ if sensor_values.shape[-1] < target:
265
+ sensor_values = F.pad(sensor_values, (0, target - sensor_values.shape[-1]))
266
+ elif sensor_values.shape[-1] > target:
267
+ sensor_values = sensor_values[..., :target]
268
+ hidden = self.source.project(sensor_values, "modality/sensor", self.config.hidden_size)
269
+ length = hidden.shape[1]
270
+ local = torch.arange(length, device=hidden.device).expand(hidden.shape[0], -1)
271
+ coords = torch.zeros(hidden.shape[0], length, 4, device=hidden.device, dtype=hidden.dtype)
272
+ coords[..., 0] = self._normalize_coords(torch.arange(length, device=hidden.device), length)
273
+ coords[..., 3] = sensor_values.float().std(dim=-1).to(hidden.dtype).clamp(max=1.0) * 2.0 - 1.0
274
+ return hidden, self._metadata(local, coords, MODALITY_SENSOR, prefix=True), (length, target)
275
+
276
+ @staticmethod
277
+ def _metadata(
278
+ local: torch.Tensor,
279
+ coords: torch.Tensor,
280
+ modality_id: int,
281
+ *,
282
+ prefix: bool,
283
+ ) -> dict[str, torch.Tensor]:
284
+ batch, length = local.shape
285
+ device = local.device
286
+ return {
287
+ "modality": torch.full((batch, length), modality_id, device=device, dtype=torch.long),
288
+ "positions": local,
289
+ "logical": local + MODALITY_POSITION_OFFSETS[modality_id],
290
+ "coords": coords,
291
+ "prefix": torch.full((batch, length), prefix, device=device, dtype=torch.bool),
292
+ "mask": torch.ones(batch, length, device=device, dtype=torch.bool),
293
+ }
294
+
295
+ def forward(
296
+ self,
297
+ *,
298
+ input_ids: torch.Tensor | None = None,
299
+ inputs_embeds: torch.Tensor | None = None,
300
+ attention_mask: torch.Tensor | None = None,
301
+ pixel_values: torch.Tensor | None = None,
302
+ audio_values: torch.Tensor | None = None,
303
+ video_values: torch.Tensor | None = None,
304
+ sensor_values: torch.Tensor | None = None,
305
+ prefix_mask: torch.Tensor | None = None,
306
+ labels: torch.Tensor | None = None,
307
+ position_start: int = 0,
308
+ ) -> DendroPackedInput:
309
+ self._batch_size(input_ids, inputs_embeds, pixel_values, audio_values, video_values, sensor_values)
310
+ # Non-text modalities form a bidirectional perceptual prefix. Text remains
311
+ # last so causal decoding can append tokens without repacking old inputs.
312
+ parts = [
313
+ ("image", MODALITY_IMAGE, self._image(pixel_values)),
314
+ ("video", MODALITY_VIDEO, self._video(video_values)),
315
+ ("audio", MODALITY_AUDIO, self._audio(audio_values)),
316
+ ("sensor", MODALITY_SENSOR, self._sensor(sensor_values)),
317
+ (
318
+ "text",
319
+ MODALITY_TEXT,
320
+ self._text(
321
+ input_ids,
322
+ inputs_embeds,
323
+ attention_mask,
324
+ position_start=position_start,
325
+ prefix_mask=prefix_mask,
326
+ ),
327
+ ),
328
+ ]
329
+ hidden_parts: list[torch.Tensor] = []
330
+ metadata: dict[str, list[torch.Tensor]] = {
331
+ "modality": [],
332
+ "positions": [],
333
+ "logical": [],
334
+ "coords": [],
335
+ "prefix": [],
336
+ "mask": [],
337
+ }
338
+ segments: list[ModalitySegment] = []
339
+ cursor = 0
340
+ text_start, text_length = 0, 0
341
+ for name, modality_id, result in parts:
342
+ if result is None:
343
+ continue
344
+ hidden, info, original_shape = result
345
+ length = int(hidden.shape[1])
346
+ # Physical sequence positions are contiguous across the packed sequence.
347
+ physical = torch.arange(cursor + position_start, cursor + position_start + length, device=hidden.device)
348
+ info["positions"] = physical.expand(hidden.shape[0], -1)
349
+ # Logical modality offsets are applied to the same absolute physical
350
+ # positions so a token keeps identical coordinates when a multimodal
351
+ # prefix is processed in one call or reused through KV cache.
352
+ info["logical"] = info["positions"] + MODALITY_POSITION_OFFSETS[modality_id]
353
+ if modality_id == MODALITY_TEXT:
354
+ denominator = max(1, self.config.max_position_embeddings - 1)
355
+ info["coords"][..., 0] = (
356
+ info["positions"].to(hidden.dtype) / denominator * 2.0 - 1.0
357
+ ).clamp(-1.0, 1.0)
358
+ hidden_parts.append(hidden)
359
+ for key in metadata:
360
+ metadata[key].append(info[key])
361
+ segments.append(ModalitySegment(name, modality_id, cursor, cursor + length, original_shape))
362
+ if modality_id == MODALITY_TEXT:
363
+ text_start, text_length = cursor, length
364
+ cursor += length
365
+ if not hidden_parts:
366
+ raise RuntimeError("No modality generated tokens")
367
+
368
+ hidden = torch.cat(hidden_parts, dim=1)
369
+ combined = {key: torch.cat(values, dim=1) for key, values in metadata.items()}
370
+ hidden = self.spatial_encoder(
371
+ hidden,
372
+ modality_ids=combined["modality"],
373
+ sequence_positions=combined["positions"],
374
+ logical_positions=combined["logical"],
375
+ coordinates=combined["coords"],
376
+ is_prefix=combined["prefix"],
377
+ )
378
+ layout = DendroModalityLayout(
379
+ modality_ids=combined["modality"],
380
+ sequence_positions=combined["positions"],
381
+ logical_positions=combined["logical"],
382
+ coordinates=combined["coords"],
383
+ is_prefix=combined["prefix"],
384
+ attention_mask=combined["mask"],
385
+ segments=tuple(segments),
386
+ text_start=text_start,
387
+ text_length=text_length,
388
+ )
389
+
390
+ aligned_labels = labels
391
+ if labels is not None:
392
+ if labels.shape[0] != hidden.shape[0]:
393
+ raise ValueError("labels batch size does not match inputs")
394
+ if labels.shape[1] == text_length and text_start > 0:
395
+ prefix_labels = torch.full(
396
+ (labels.shape[0], text_start),
397
+ -100,
398
+ device=labels.device,
399
+ dtype=labels.dtype,
400
+ )
401
+ aligned_labels = torch.cat([prefix_labels, labels], dim=1)
402
+ elif labels.shape[1] != hidden.shape[1]:
403
+ raise ValueError(
404
+ f"labels length {labels.shape[1]} must equal text length {text_length} "
405
+ f"or packed length {hidden.shape[1]}"
406
+ )
407
+ return DendroPackedInput(hidden_states=hidden, layout=layout, aligned_labels=aligned_labels)
modeling_dendro_omni.py ADDED
@@ -0,0 +1,1929 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Hugging Face compatible unified Dendro Omni model."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import math
6
+ from dataclasses import dataclass
7
+ from pathlib import Path
8
+ from typing import Any, Callable, Mapping, Sequence
9
+
10
+ import torch
11
+ from torch import nn
12
+ from torch.nn import functional as F
13
+ from torch.utils.checkpoint import checkpoint as activation_checkpoint
14
+
15
+ from ._source_bound import SourceBoundModule
16
+ from .cache import DendroKVCache
17
+ from .cell import DendroCellState, DendroRecurrentCell
18
+ from .configuration_dendro_omni import DendroOmniConfig
19
+ from .hf_compat import (
20
+ CausalLMOutputWithPast,
21
+ DendroGenerationConfig,
22
+ GenerationMixin,
23
+ ModelOutput,
24
+ PreTrainedModel,
25
+ TRANSFORMERS_AVAILABLE,
26
+ )
27
+ from .modalities import DendroModalityLayout, DendroOmniInputProjector
28
+ from .source import DendroSourceLayer
29
+ from .transplant import DendroExactTransplantCore, DendroTransplantCache
30
+
31
+
32
+ @dataclass(slots=True)
33
+ class DendroAwakeningCache:
34
+ """Incremental donor and native state without duplicate model weights."""
35
+
36
+ donor: DendroTransplantCache
37
+ native: DendroKVCache
38
+
39
+ @property
40
+ def total_tokens_seen(self) -> int:
41
+ return self.donor.total_tokens_seen
42
+
43
+ def get_seq_length(self, layer_idx: int = 0) -> int:
44
+ del layer_idx
45
+ return self.donor.get_seq_length()
46
+
47
+ def reorder_cache(self, indices: torch.Tensor) -> "DendroAwakeningCache":
48
+ self.donor.reorder_cache(indices)
49
+ self.native.reorder_cache(indices)
50
+ return self
51
+
52
+
53
+ @dataclass
54
+ class DendroBaseModelOutputWithPast(ModelOutput):
55
+ last_hidden_state: torch.Tensor | None = None
56
+ past_key_values: DendroKVCache | DendroTransplantCache | DendroAwakeningCache | tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None
57
+ hidden_states: tuple[torch.Tensor, ...] | None = None
58
+ attentions: tuple[torch.Tensor, ...] | None = None
59
+ reasoning_state: dict[str, Any] | None = None
60
+ modality_layout: dict[str, Any] | None = None
61
+ dendro_states: tuple[DendroCellState, ...] | None = None
62
+ aligned_labels: torch.Tensor | None = None
63
+
64
+
65
+ @dataclass
66
+ class DendroCausalLMOutput(CausalLMOutputWithPast):
67
+ reasoning_state: dict[str, Any] | None = None
68
+ modality_layout: dict[str, Any] | None = None
69
+ dendro_states: tuple[DendroCellState, ...] | None = None
70
+
71
+
72
+ @dataclass(slots=True)
73
+ class DendroAnswerGenerationOutput:
74
+ """Answer-only result from hidden deliberative generation.
75
+
76
+ Deliberation token IDs are intentionally absent. The answer mask identifies
77
+ real tokens in the padded batch and is directly usable for tokenizer decoding.
78
+ """
79
+
80
+ answer_token_ids: torch.Tensor
81
+ answer_attention_mask: torch.Tensor
82
+ reasoning_token_counts: tuple[int, ...]
83
+ reasoning_minimum_tokens: tuple[int, ...]
84
+ reasoning_token_budgets: tuple[int, ...]
85
+ reasoning_completed: tuple[bool, ...]
86
+ reasoning_forced_close: tuple[bool, ...]
87
+ answer_token_counts: tuple[int, ...]
88
+ answer_eos_reached: tuple[bool, ...]
89
+ answer_reasoning_filtered: tuple[bool, ...]
90
+ finish_reasons: tuple[str, ...]
91
+
92
+ @property
93
+ def sequences(self) -> torch.Tensor:
94
+ """Answer-only tensor alias for generation consumers."""
95
+
96
+ return self.answer_token_ids
97
+
98
+ def batch_decode(
99
+ self,
100
+ tokenizer: Any,
101
+ *,
102
+ skip_special_tokens: bool = True,
103
+ ) -> list[str]:
104
+ rows = [
105
+ row[mask].detach().cpu().tolist()
106
+ for row, mask in zip(self.answer_token_ids, self.answer_attention_mask)
107
+ ]
108
+ if hasattr(tokenizer, "batch_decode"):
109
+ try:
110
+ return list(
111
+ tokenizer.batch_decode(
112
+ rows,
113
+ skip_special_tokens=skip_special_tokens,
114
+ clean_up_tokenization_spaces=False,
115
+ )
116
+ )
117
+ except TypeError:
118
+ return list(
119
+ tokenizer.batch_decode(
120
+ rows, skip_special_tokens=skip_special_tokens
121
+ )
122
+ )
123
+ return [
124
+ str(tokenizer.decode(row, skip_special_tokens=skip_special_tokens))
125
+ for row in rows
126
+ ]
127
+
128
+
129
+ class DendroSourceEmbedding(SourceBoundModule):
130
+ """HF-facing embedding proxy with no private parameters."""
131
+
132
+ def __init__(self, config: DendroOmniConfig, source: DendroSourceLayer) -> None:
133
+ super().__init__(source)
134
+ self.num_embeddings = config.vocab_size
135
+ self.embedding_dim = config.hidden_size
136
+
137
+ @property
138
+ def weight(self) -> torch.Tensor:
139
+ return self.source.primitive("token/embedding", (self.num_embeddings, self.embedding_dim))
140
+
141
+ def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
142
+ return F.embedding(input_ids, self.weight)
143
+
144
+
145
+ class DendroSourceLMHead(SourceBoundModule):
146
+ """HF-facing tied output proxy over the same token source view."""
147
+
148
+ def __init__(self, config: DendroOmniConfig, source: DendroSourceLayer) -> None:
149
+ super().__init__(source)
150
+ self.in_features = config.hidden_size
151
+ self.out_features = config.vocab_size
152
+
153
+ @property
154
+ def weight(self) -> torch.Tensor:
155
+ return self.source.primitive("token/embedding", (self.out_features, self.in_features))
156
+
157
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
158
+ return self.source.tied_logits(hidden_states, vocab_size=self.out_features)
159
+
160
+
161
+ class DendroPreTrainedModel(PreTrainedModel):
162
+ config_class = DendroOmniConfig
163
+ generation_config_class = DendroGenerationConfig
164
+ base_model_prefix = "model"
165
+ main_input_name = "input_ids"
166
+ supports_gradient_checkpointing = False
167
+ _supports_sdpa = True
168
+ _supports_cache_class = True
169
+ _supports_static_cache = True
170
+ _no_split_modules = ["DendroRecurrentCell", "DendroSourceLayer"]
171
+ _skip_keys_device_placement = ["past_key_values"]
172
+
173
+ def _init_weights(self, module: nn.Module) -> None:
174
+ # The source tensor initializes itself exactly once. Parameterless views must
175
+ # not trigger Hugging Face's per-module initializer.
176
+ del module
177
+
178
+ @classmethod
179
+ def _supports_default_dynamic_cache(cls) -> bool:
180
+ """Prevent GenerationMixin from injecting a generic HF cache.
181
+
182
+ Dendro must instantiate :class:`DendroKVCache` itself because each virtual
183
+ recurrent depth owns an activation-only K/V stream plus workspace, memory,
184
+ plasticity and associative state. Returning ``False`` is the current
185
+ Transformers extension point for architectures with a custom cache.
186
+ """
187
+
188
+ return False
189
+
190
+
191
+ if issubclass(PreTrainedModel, GenerationMixin):
192
+ class _DendroGenerationBase(DendroPreTrainedModel):
193
+ pass
194
+ else:
195
+ class _DendroGenerationBase(DendroPreTrainedModel, GenerationMixin):
196
+ pass
197
+
198
+
199
+ class DendroOmniModel(DendroPreTrainedModel):
200
+ """The recurrent backbone: one source layer and one physical cell."""
201
+
202
+ def __init__(self, config: DendroOmniConfig) -> None:
203
+ super().__init__(config)
204
+ self.source_layer = DendroSourceLayer(
205
+ config.source_size,
206
+ low_bit=config.low_bit,
207
+ ternary_threshold=config.ternary_threshold,
208
+ init_std=config.init_std,
209
+ cache_derived_views=config.cache_derived_views,
210
+ cache_indices=config.cache_source_indices,
211
+ logical_region_start=config.source_logical_region_start,
212
+ logical_region_size=config.source_logical_region_size,
213
+ tensor_map=config.transplant_tensor_map,
214
+ aliases=config.transplant_aliases,
215
+ exact_tied_logits=config.transplant_exact_tied_logits,
216
+ freeze_named_tensors=config.transplant_freeze_donor_bank,
217
+ )
218
+ self.transplant_core = (
219
+ DendroExactTransplantCore(config, self.source_layer)
220
+ if config.transplant_mode in {"exact", "awakening"}
221
+ else None
222
+ )
223
+ if self.transplant_core is not None:
224
+ missing = self.transplant_core.validate_required_tensors()
225
+ if missing:
226
+ preview = ", ".join(missing[:8])
227
+ raise ValueError(
228
+ f"transplant checkpoint is missing {len(missing)} required donor tensors: {preview}"
229
+ )
230
+ self.input_projector = DendroOmniInputProjector(config, self.source_layer)
231
+ self.cell = DendroRecurrentCell(config, self.source_layer)
232
+ self._input_embedding_proxy = DendroSourceEmbedding(config, self.source_layer)
233
+ self.post_init()
234
+
235
+ def get_input_embeddings(self) -> DendroSourceEmbedding:
236
+ return self._input_embedding_proxy
237
+
238
+ def set_input_embeddings(self, value: nn.Module) -> None:
239
+ weight = getattr(value, "weight", None)
240
+ if weight is None or tuple(weight.shape) != (self.config.vocab_size, self.config.hidden_size):
241
+ raise ValueError("input embedding must expose weight [vocab_size, hidden_size]")
242
+ self.source_layer.write_primitive("token/embedding", weight)
243
+
244
+ def _new_cache(self) -> DendroKVCache:
245
+ return DendroKVCache(
246
+ max_depth=self.config.max_total_recurrent_steps,
247
+ implementation=self.config.cache_implementation,
248
+ max_cache_length=self.config.max_cache_length,
249
+ sliding_window=self.config.sliding_window,
250
+ offload_device=self.config.cache_offload_device,
251
+ detach_runtime_state=self.config.detach_runtime_state,
252
+ )
253
+
254
+ @staticmethod
255
+ def _mean_metric(states: list[DendroCellState], name: str) -> float:
256
+ values = [getattr(state, name).detach().float().mean() for state in states]
257
+ return float(torch.stack(values).mean().cpu().item()) if values else 0.0
258
+
259
+ def forward(
260
+ self,
261
+ input_ids: torch.Tensor | None = None,
262
+ attention_mask: torch.Tensor | None = None,
263
+ position_ids: torch.Tensor | None = None,
264
+ past_key_values: DendroKVCache | DendroTransplantCache | DendroAwakeningCache | tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None,
265
+ inputs_embeds: torch.Tensor | None = None,
266
+ pixel_values: torch.Tensor | None = None,
267
+ pixel_values_videos: torch.Tensor | None = None,
268
+ image_grid_thw: torch.Tensor | None = None,
269
+ video_grid_thw: torch.Tensor | None = None,
270
+ mm_token_type_ids: torch.Tensor | None = None,
271
+ audio_values: torch.Tensor | None = None,
272
+ video_values: torch.Tensor | None = None,
273
+ sensor_values: torch.Tensor | None = None,
274
+ prefix_mask: torch.Tensor | None = None,
275
+ labels: torch.Tensor | None = None,
276
+ reasoning_effort: str | None = None,
277
+ reasoning_budget: int | None = None,
278
+ use_cache: bool | None = None,
279
+ output_attentions: bool | None = None,
280
+ output_hidden_states: bool | None = None,
281
+ return_dict: bool | None = None,
282
+ collect_reasoning_diagnostics: bool = True,
283
+ cache_position: torch.Tensor | None = None,
284
+ **extra_kwargs: Any,
285
+ ) -> DendroBaseModelOutputWithPast | tuple[Any, ...]:
286
+ native_training_only = bool(extra_kwargs.pop("native_training_only", False))
287
+ training_recurrent_depth = extra_kwargs.pop("training_recurrent_depth", None)
288
+ if extra_kwargs:
289
+ unknown = ", ".join(sorted(extra_kwargs))
290
+ raise TypeError(f"unexpected Dendro model arguments: {unknown}")
291
+ del cache_position # Positions are represented directly in the packed layout.
292
+ use_cache = self.config.use_cache if use_cache is None else bool(use_cache)
293
+ output_attentions = self.config.output_attentions if output_attentions is None else bool(output_attentions)
294
+ output_hidden_states = (
295
+ self.config.output_hidden_states if output_hidden_states is None else bool(output_hidden_states)
296
+ )
297
+ return_dict = self.config.use_return_dict if return_dict is None else bool(return_dict)
298
+ collect_reasoning_diagnostics = bool(collect_reasoning_diagnostics)
299
+ self.source_layer.begin_forward_cache()
300
+
301
+ explicit_reasoning_budget = None
302
+ if reasoning_budget is not None:
303
+ explicit_reasoning_budget = self.config.resolve_reasoning_budget(
304
+ reasoning_effort, reasoning_budget
305
+ )
306
+
307
+ transplant_mode = self.config.transplant_mode
308
+ transplant_blend = self.config.transplant_blend_for_effort(reasoning_effort)
309
+ if self.transplant_core is not None and (transplant_mode == "exact" or transplant_blend <= 0.0):
310
+ explicit_reasoning_budget = 0
311
+ if past_key_values is not None and not isinstance(past_key_values, DendroTransplantCache):
312
+ raise ValueError(
313
+ "exact transplant caching requires DendroTransplantCache"
314
+ )
315
+ if any(value is not None for value in (audio_values, video_values, sensor_values, prefix_mask)):
316
+ raise ValueError(
317
+ "Qwen3.5 exact transplant mode accepts Qwen text/image/video processor tensors; "
318
+ "Dendro-native audio/sensor inputs become available in awakening mode"
319
+ )
320
+ if input_ids is None and inputs_embeds is None:
321
+ raise ValueError("exact transplant mode requires input_ids or inputs_embeds")
322
+ if input_ids is not None:
323
+ transplant_output = self.transplant_core.multimodal_forward(
324
+ input_ids=input_ids,
325
+ attention_mask=attention_mask,
326
+ position_ids=position_ids,
327
+ pixel_values=pixel_values,
328
+ pixel_values_videos=pixel_values_videos,
329
+ image_grid_thw=image_grid_thw,
330
+ video_grid_thw=video_grid_thw,
331
+ mm_token_type_ids=mm_token_type_ids,
332
+ cache=past_key_values,
333
+ use_cache=use_cache,
334
+ )
335
+ else:
336
+ if any(value is not None for value in (pixel_values, pixel_values_videos)):
337
+ raise ValueError("multimodal exact transplant mode requires input_ids for placeholder placement")
338
+ transplant_output = self.transplant_core.text_forward(
339
+ inputs_embeds=inputs_embeds,
340
+ attention_mask=attention_mask,
341
+ position_ids=position_ids,
342
+ cache=past_key_values,
343
+ use_cache=use_cache,
344
+ )
345
+ reasoning_state = {
346
+ "transplant_mode": "exact",
347
+ "source_model": self.config.transplant_source_model,
348
+ "function_path": "dendro_native_exact_equivalence",
349
+ "dendro_extensions_active": False,
350
+ "cache": None,
351
+ }
352
+ result = DendroBaseModelOutputWithPast(
353
+ last_hidden_state=transplant_output.last_hidden_state,
354
+ past_key_values=transplant_output.cache if use_cache else None,
355
+ hidden_states=(transplant_output.last_hidden_state,) if output_hidden_states else None,
356
+ attentions=None,
357
+ reasoning_state=reasoning_state,
358
+ modality_layout={"mode": "qwen3.5_exact_transplant"},
359
+ dendro_states=tuple(),
360
+ aligned_labels=labels,
361
+ )
362
+ return result if return_dict else result.to_tuple()
363
+
364
+ if self.transplant_core is not None and transplant_mode == "awakening" and transplant_blend > 0.0:
365
+ # Preserve the donor function while admitting the original Dendro recurrent body.
366
+ # Qwen image/video features are first formed by the exact donor coordinate path,
367
+ # then the resulting fused embeddings are routed through Dendro's own cell. Native
368
+ # audio/video/sensor prefixes can be added around the same text-aligned sequence.
369
+ if past_key_values is not None and not isinstance(past_key_values, DendroAwakeningCache):
370
+ raise ValueError("awakening transplant caching requires DendroAwakeningCache")
371
+ awakening_cache = past_key_values if isinstance(past_key_values, DendroAwakeningCache) else None
372
+ donor_cache = None if awakening_cache is None else awakening_cache.donor
373
+ native_cache = None if awakening_cache is None else awakening_cache.native
374
+ if native_training_only:
375
+ if not self.training:
376
+ raise ValueError("native_training_only is restricted to model.train()")
377
+ if use_cache or past_key_values is not None:
378
+ raise ValueError("native_training_only does not retain an inference cache")
379
+ if any(
380
+ value is not None
381
+ for value in (
382
+ pixel_values,
383
+ pixel_values_videos,
384
+ image_grid_thw,
385
+ video_grid_thw,
386
+ mm_token_type_ids,
387
+ audio_values,
388
+ video_values,
389
+ sensor_values,
390
+ )
391
+ ):
392
+ raise ValueError("native_training_only currently accepts text batches only")
393
+ if inputs_embeds is None:
394
+ if input_ids is None:
395
+ raise ValueError("native_training_only requires input_ids or inputs_embeds")
396
+ inputs_embeds = self.transplant_core.token_embedding(input_ids)
397
+ original_mode = self.config.transplant_mode
398
+ self.config.transplant_mode = "off"
399
+ try:
400
+ native_output = self.forward(
401
+ input_ids=None,
402
+ attention_mask=attention_mask,
403
+ position_ids=position_ids,
404
+ inputs_embeds=inputs_embeds,
405
+ prefix_mask=prefix_mask,
406
+ labels=labels,
407
+ reasoning_effort=reasoning_effort,
408
+ reasoning_budget=reasoning_budget,
409
+ training_recurrent_depth=training_recurrent_depth,
410
+ use_cache=False,
411
+ output_attentions=output_attentions,
412
+ output_hidden_states=output_hidden_states,
413
+ return_dict=True,
414
+ collect_reasoning_diagnostics=collect_reasoning_diagnostics,
415
+ )
416
+ finally:
417
+ self.config.transplant_mode = original_mode
418
+ native_output.reasoning_state = dict(native_output.reasoning_state or {})
419
+ native_output.reasoning_state.update(
420
+ {
421
+ "transplant_mode": "awakening_native_training",
422
+ "final_inference_blend": transplant_blend,
423
+ "donor_forward_skipped": True,
424
+ }
425
+ )
426
+ return native_output if return_dict else native_output.to_tuple()
427
+ if input_ids is not None:
428
+ exact_output = self.transplant_core.multimodal_forward(
429
+ input_ids=input_ids,
430
+ attention_mask=attention_mask,
431
+ position_ids=position_ids,
432
+ pixel_values=pixel_values,
433
+ pixel_values_videos=pixel_values_videos,
434
+ image_grid_thw=image_grid_thw,
435
+ video_grid_thw=video_grid_thw,
436
+ mm_token_type_ids=mm_token_type_ids,
437
+ cache=donor_cache,
438
+ use_cache=use_cache,
439
+ )
440
+ else:
441
+ if inputs_embeds is None:
442
+ raise ValueError("awakening transplant mode requires input_ids or inputs_embeds")
443
+ if pixel_values is not None or pixel_values_videos is not None:
444
+ raise ValueError("Qwen image/video awakening requires input_ids for placeholder placement")
445
+ exact_output = self.transplant_core.text_forward(
446
+ inputs_embeds=inputs_embeds,
447
+ attention_mask=attention_mask,
448
+ position_ids=position_ids,
449
+ cache=donor_cache,
450
+ use_cache=use_cache,
451
+ )
452
+
453
+ native_prefix_mask = prefix_mask
454
+ if native_prefix_mask is None and mm_token_type_ids is not None:
455
+ native_prefix_mask = mm_token_type_ids.to(dtype=torch.bool)
456
+
457
+ original_mode = self.config.transplant_mode
458
+ self.config.transplant_mode = "off"
459
+ try:
460
+ native_output = self.forward(
461
+ input_ids=None,
462
+ attention_mask=attention_mask,
463
+ # The awakening branches must see the same absolute text
464
+ # coordinates. This is especially important for the
465
+ # trainer's randomized high-position augmentation; using
466
+ # donor offsets with native zero-based positions would blend
467
+ # representations for different coordinates.
468
+ position_ids=position_ids,
469
+ past_key_values=native_cache,
470
+ inputs_embeds=exact_output.input_embeddings,
471
+ pixel_values=None,
472
+ audio_values=audio_values,
473
+ video_values=video_values,
474
+ sensor_values=sensor_values,
475
+ prefix_mask=native_prefix_mask,
476
+ labels=labels,
477
+ reasoning_effort=reasoning_effort,
478
+ reasoning_budget=reasoning_budget,
479
+ use_cache=use_cache,
480
+ output_attentions=output_attentions,
481
+ output_hidden_states=output_hidden_states,
482
+ return_dict=True,
483
+ collect_reasoning_diagnostics=collect_reasoning_diagnostics,
484
+ )
485
+ finally:
486
+ self.config.transplant_mode = original_mode
487
+
488
+ native_hidden = native_output.last_hidden_state
489
+ if native_hidden is None:
490
+ raise RuntimeError("Dendro awakening path produced no hidden state")
491
+ layout = native_output.modality_layout or {}
492
+ text_start = int(layout.get("text_start", 0))
493
+ text_length = int(layout.get("text_length", native_hidden.shape[1]))
494
+ if text_length != exact_output.last_hidden_state.shape[1]:
495
+ raise RuntimeError(
496
+ "Dendro awakening text segment is not aligned with the exact donor sequence: "
497
+ f"native={text_length}, exact={exact_output.last_hidden_state.shape[1]}"
498
+ )
499
+ blended = transplant_blend * native_hidden
500
+ blended[:, text_start : text_start + text_length] = (
501
+ blended[:, text_start : text_start + text_length]
502
+ + exact_output.last_hidden_state
503
+ )
504
+ native_output.last_hidden_state = blended
505
+ native_output.reasoning_state = dict(native_output.reasoning_state or {})
506
+ native_output.reasoning_state.update(
507
+ {
508
+ "transplant_mode": "awakening",
509
+ "transplant_native_blend": transplant_blend,
510
+ "source_model": self.config.transplant_source_model,
511
+ }
512
+ )
513
+ if use_cache:
514
+ if exact_output.cache is None or not isinstance(native_output.past_key_values, DendroKVCache):
515
+ raise RuntimeError("awakening cache initialization failed")
516
+ native_output.past_key_values = DendroAwakeningCache(
517
+ donor=exact_output.cache,
518
+ native=native_output.past_key_values,
519
+ )
520
+ return native_output if return_dict else native_output.to_tuple()
521
+
522
+ if past_key_values is not None and not isinstance(past_key_values, DendroKVCache):
523
+ past_key_values = DendroKVCache.from_legacy_cache(
524
+ past_key_values,
525
+ max_depth=self.config.max_total_recurrent_steps,
526
+ implementation=self.config.cache_implementation,
527
+ )
528
+ cache = past_key_values
529
+ if use_cache and cache is None:
530
+ cache = self._new_cache()
531
+ if cache is not None and cache.get_seq_length() > 0:
532
+ if any(value is not None for value in (pixel_values, audio_values, video_values, sensor_values)):
533
+ raise ValueError("Multimodal prefix tensors may only be supplied while the cache is empty")
534
+ position_start = cache.total_tokens_seen
535
+ else:
536
+ position_start = 0
537
+
538
+ packed = self.input_projector(
539
+ input_ids=input_ids,
540
+ inputs_embeds=inputs_embeds,
541
+ attention_mask=attention_mask,
542
+ pixel_values=pixel_values,
543
+ audio_values=audio_values,
544
+ video_values=video_values,
545
+ sensor_values=sensor_values,
546
+ prefix_mask=prefix_mask,
547
+ labels=labels,
548
+ position_start=position_start,
549
+ )
550
+ if position_ids is not None:
551
+ if len(packed.layout.segments) != 1 or packed.layout.segments[0].name != "text":
552
+ raise ValueError("Explicit position_ids are supported only for text-only inputs")
553
+ if position_ids.shape != packed.layout.sequence_positions.shape:
554
+ raise ValueError("position_ids shape must match input_ids")
555
+ packed.layout.sequence_positions = position_ids.to(packed.hidden_states.device)
556
+
557
+ effort_name, effort_id, effort_level = self.config.reasoning_effort_profile(
558
+ reasoning_effort
559
+ )
560
+ extra_requested = self.config.resolve_reasoning_budget(
561
+ effort_name, reasoning_budget
562
+ )
563
+ verifier_passes = self.config.reasoning_verifier_passes if extra_requested >= 3 else 0
564
+ base_depth = self.config.num_hidden_layers
565
+ if training_recurrent_depth is not None:
566
+ if not self.training:
567
+ raise ValueError("training_recurrent_depth is restricted to model.train()")
568
+ base_depth = int(training_recurrent_depth)
569
+ if not 1 <= base_depth <= self.config.num_hidden_layers:
570
+ raise ValueError("training_recurrent_depth must be within the configured recurrent depth")
571
+ # Legacy shallow training remains unchanged unless the caller opts in
572
+ # with an explicit reasoning budget. This lets a future curriculum
573
+ # train reasoning/verification phases without requiring full base depth.
574
+ if reasoning_budget is None:
575
+ extra_requested = 0
576
+ verifier_passes = 0
577
+ requested_depth = base_depth + extra_requested + verifier_passes
578
+ reasoning_signature = (
579
+ effort_name,
580
+ effort_id,
581
+ round(effort_level, 8),
582
+ extra_requested,
583
+ verifier_passes,
584
+ )
585
+ if cache is not None:
586
+ cache.ensure_recurrent_depth(requested_depth)
587
+ previous_signature = cache.reasoning_state.get("reasoning_signature")
588
+ if (
589
+ cache.get_seq_length() > 0
590
+ and previous_signature is not None
591
+ and tuple(previous_signature) != reasoning_signature
592
+ ):
593
+ raise RuntimeError(
594
+ "reasoning effort/budget changed while reusing a populated cache; "
595
+ "start a fresh cache or keep the same reasoning controls"
596
+ )
597
+
598
+ hidden = packed.hidden_states
599
+ hidden_history: list[torch.Tensor] = [hidden] if output_hidden_states else []
600
+ attention_history: list[torch.Tensor] = []
601
+ states: list[DendroCellState] = []
602
+
603
+ checkpoint_recurrence = (
604
+ self.training
605
+ and torch.is_grad_enabled()
606
+ and bool(getattr(self.config, "training_recurrent_checkpointing", False))
607
+ and not use_cache
608
+ )
609
+
610
+ def apply_cell(
611
+ current_hidden: torch.Tensor,
612
+ *,
613
+ current_depth: int,
614
+ current_phase: str,
615
+ current_progress: float,
616
+ current_remaining: float,
617
+ ):
618
+ def cell_forward(value: torch.Tensor):
619
+ return self.cell(
620
+ value,
621
+ layout=packed.layout,
622
+ depth_idx=current_depth,
623
+ phase=current_phase,
624
+ effort_id=effort_id,
625
+ effort_level=effort_level,
626
+ phase_progress=current_progress,
627
+ remaining_budget_fraction=current_remaining,
628
+ cache=cache,
629
+ use_cache=use_cache,
630
+ output_attentions=output_attentions,
631
+ )
632
+
633
+ if not checkpoint_recurrence:
634
+ return cell_forward(current_hidden)
635
+ return activation_checkpoint(
636
+ cell_forward,
637
+ current_hidden,
638
+ use_reentrant=False,
639
+ preserve_rng_state=True,
640
+ )
641
+
642
+ depth_idx = 0
643
+ for step in range(base_depth):
644
+ output = apply_cell(
645
+ hidden,
646
+ current_depth=depth_idx,
647
+ current_phase="base",
648
+ current_progress=(step + 1) / max(1, base_depth),
649
+ current_remaining=(requested_depth - depth_idx - 1)
650
+ / max(1, requested_depth),
651
+ )
652
+ hidden = output.hidden_states
653
+ states.append(output.state)
654
+ if output_hidden_states:
655
+ hidden_history.append(hidden)
656
+ if output_attentions and output.attention_weights is not None:
657
+ attention_history.append(output.attention_weights)
658
+ depth_idx += 1
659
+
660
+ extra_executed = 0
661
+ refinement_scale = float(self.config.reasoning_refinement_scale)
662
+ adaptive_allowed = (
663
+ self.config.adaptive_reasoning
664
+ and not use_cache
665
+ and reasoning_budget is None
666
+ )
667
+ for step in range(extra_requested):
668
+ output = apply_cell(
669
+ hidden,
670
+ current_depth=depth_idx,
671
+ current_phase="reasoning",
672
+ current_progress=(step + 1) / max(1, extra_requested),
673
+ current_remaining=(requested_depth - depth_idx - 1)
674
+ / max(1, requested_depth),
675
+ )
676
+ if refinement_scale >= 1.0:
677
+ hidden = output.hidden_states
678
+ elif refinement_scale > 0.0:
679
+ hidden = hidden + refinement_scale * (output.hidden_states - hidden)
680
+ states.append(output.state)
681
+ extra_executed += 1
682
+ if output_hidden_states:
683
+ hidden_history.append(hidden)
684
+ if output_attentions and output.attention_weights is not None:
685
+ attention_history.append(output.attention_weights)
686
+ depth_idx += 1
687
+ if adaptive_allowed and extra_executed >= self.config.reasoning_min_steps:
688
+ ready = float(output.state.readiness.detach().float().mean().cpu().item())
689
+ contradiction = float(output.state.contradiction.detach().float().mean().cpu().item())
690
+ if ready * (1.0 - contradiction) >= self.config.reasoning_stop_threshold:
691
+ break
692
+
693
+ verifier_executed = 0
694
+ # Cached generation uses a fixed depth so every recurrent K/V stream remains
695
+ # complete. In non-cached analysis, verifier passes follow actual reasoning.
696
+ passes = verifier_passes if use_cache else (verifier_passes if extra_executed else 0)
697
+ for step in range(passes):
698
+ output = apply_cell(
699
+ hidden,
700
+ current_depth=depth_idx,
701
+ current_phase="verification",
702
+ current_progress=(step + 1) / max(1, passes),
703
+ current_remaining=(requested_depth - depth_idx - 1)
704
+ / max(1, requested_depth),
705
+ )
706
+ if refinement_scale >= 1.0:
707
+ hidden = output.hidden_states
708
+ elif refinement_scale > 0.0:
709
+ hidden = hidden + refinement_scale * (output.hidden_states - hidden)
710
+ states.append(output.state)
711
+ verifier_executed += 1
712
+ if output_hidden_states:
713
+ hidden_history.append(hidden)
714
+ if output_attentions and output.attention_weights is not None:
715
+ attention_history.append(output.attention_weights)
716
+ depth_idx += 1
717
+
718
+ hidden = self.source_layer.rms_norm(hidden, "model/final_norm", eps=self.config.layer_norm_eps)
719
+ if output_hidden_states:
720
+ hidden_history.append(hidden)
721
+
722
+ reasoning_state: dict[str, Any] = {
723
+ "effort": effort_name,
724
+ "effort_id": effort_id,
725
+ "effort_level": effort_level,
726
+ "reasoning_budget": extra_requested,
727
+ "reasoning_budget_override": reasoning_budget is not None,
728
+ "base_recurrent_steps": self.config.num_hidden_layers,
729
+ "reasoning_steps_requested": extra_requested,
730
+ "reasoning_steps_executed": extra_executed,
731
+ "verifier_steps_executed": verifier_executed,
732
+ "total_recurrent_steps": depth_idx,
733
+ "adaptive_early_stop": bool(adaptive_allowed and extra_executed < extra_requested),
734
+ "reasoning_route_prior_strength": float(
735
+ self.config.reasoning_route_prior_strength
736
+ ),
737
+ "reasoning_refinement_scale": refinement_scale,
738
+ "diagnostics_collected": collect_reasoning_diagnostics,
739
+ "reasoning_signature": list(reasoning_signature),
740
+ }
741
+ if collect_reasoning_diagnostics:
742
+ final_state = states[-1]
743
+ readiness = float(final_state.readiness.detach().float().mean().cpu().item())
744
+ contradiction = float(final_state.contradiction.detach().float().mean().cpu().item())
745
+ route_distribution = (
746
+ final_state.route_probs.detach().float().mean(dim=(0, 1)).cpu()
747
+ )
748
+ expert_distribution = (
749
+ final_state.expert_probs.detach().float().mean(dim=(0, 1)).cpu()
750
+ )
751
+ route_entropy = float(
752
+ (
753
+ -route_distribution
754
+ * route_distribution.clamp_min(1e-8).log()
755
+ ).sum().item()
756
+ )
757
+ reasoning_state.update(
758
+ {
759
+ "readiness": readiness,
760
+ "contradiction": contradiction,
761
+ "confidence": readiness * (1.0 - contradiction),
762
+ "mean_coherence": self._mean_metric(states, "coherence"),
763
+ "mean_entropy_pressure": self._mean_metric(states, "entropy_pressure"),
764
+ "mean_salience": self._mean_metric(states, "salience"),
765
+ "route_distribution": route_distribution.tolist(),
766
+ "route_entropy": route_entropy,
767
+ "expert_distribution": expert_distribution.tolist(),
768
+ "cache": cache.summary() if cache is not None else None,
769
+ }
770
+ )
771
+ if cache is not None:
772
+ cache.reasoning_state = reasoning_state
773
+
774
+ result = DendroBaseModelOutputWithPast(
775
+ last_hidden_state=hidden,
776
+ past_key_values=cache if use_cache else None,
777
+ hidden_states=tuple(hidden_history) if output_hidden_states else None,
778
+ attentions=tuple(attention_history) if output_attentions else None,
779
+ reasoning_state=reasoning_state,
780
+ modality_layout=packed.layout.to_dict(),
781
+ dendro_states=tuple(states),
782
+ aligned_labels=packed.aligned_labels,
783
+ )
784
+ if return_dict:
785
+ return result
786
+ return result.to_tuple()
787
+
788
+
789
+ class DendroForCausalLM(_DendroGenerationBase):
790
+ """Causal LM with tied logits derived from the same single source tensor."""
791
+
792
+ _tied_weights_keys = {}
793
+
794
+ def __init__(self, config: DendroOmniConfig) -> None:
795
+ super().__init__(config)
796
+ self.model = DendroOmniModel(config)
797
+ self._output_embedding_proxy = DendroSourceLMHead(
798
+ config,
799
+ self.model.source_layer,
800
+ )
801
+ self.post_init()
802
+
803
+ @property
804
+ def source_layer(self) -> DendroSourceLayer:
805
+ return self.model.source_layer
806
+
807
+ def get_input_embeddings(self) -> DendroSourceEmbedding:
808
+ return self.model.get_input_embeddings()
809
+
810
+ def set_input_embeddings(self, value: nn.Module) -> None:
811
+ self.model.set_input_embeddings(value)
812
+
813
+ def get_output_embeddings(self) -> DendroSourceLMHead:
814
+ return self._output_embedding_proxy
815
+
816
+ def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
817
+ weight = getattr(new_embeddings, "weight", None)
818
+ if weight is None or tuple(weight.shape) != (self.config.vocab_size, self.config.hidden_size):
819
+ raise ValueError("output embedding must expose weight [vocab_size, hidden_size]")
820
+ self.source_layer.write_primitive("token/embedding", weight)
821
+
822
+ def tie_weights(
823
+ self,
824
+ missing_keys: set[str] | None = None,
825
+ recompute_mapping: bool = True,
826
+ ) -> None:
827
+ del missing_keys, recompute_mapping
828
+ return None
829
+
830
+ def resize_token_embeddings(
831
+ self,
832
+ new_num_tokens: int | None = None,
833
+ pad_to_multiple_of: int | None = None,
834
+ mean_resizing: bool = True,
835
+ ) -> DendroSourceEmbedding:
836
+ del mean_resizing
837
+ if new_num_tokens is None:
838
+ return self.get_input_embeddings()
839
+ if pad_to_multiple_of:
840
+ new_num_tokens = math.ceil(new_num_tokens / pad_to_multiple_of) * pad_to_multiple_of
841
+ if new_num_tokens < self.config.byte_offset + 256:
842
+ raise ValueError("The byte tokenizer requires at least byte_offset + 256 tokens")
843
+ old_weight = self.get_input_embeddings().weight.detach().clone()
844
+ self.source_layer.forget_primitive_shape("token")
845
+ self.config.vocab_size = int(new_num_tokens)
846
+ self.model._input_embedding_proxy.num_embeddings = int(new_num_tokens)
847
+ self._output_embedding_proxy.out_features = int(new_num_tokens)
848
+ resized = torch.empty(
849
+ new_num_tokens,
850
+ self.config.hidden_size,
851
+ device=old_weight.device,
852
+ dtype=old_weight.dtype,
853
+ )
854
+ nn.init.normal_(resized, mean=0.0, std=self.config.init_std)
855
+ count = min(old_weight.shape[0], new_num_tokens)
856
+ resized[:count].copy_(old_weight[:count])
857
+ self.source_layer.write_primitive("token/embedding", resized)
858
+ return self.get_input_embeddings()
859
+
860
+ def forward(
861
+ self,
862
+ input_ids: torch.Tensor | None = None,
863
+ attention_mask: torch.Tensor | None = None,
864
+ position_ids: torch.Tensor | None = None,
865
+ past_key_values: DendroKVCache | DendroTransplantCache | DendroAwakeningCache | tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None,
866
+ inputs_embeds: torch.Tensor | None = None,
867
+ labels: torch.Tensor | None = None,
868
+ pixel_values: torch.Tensor | None = None,
869
+ pixel_values_videos: torch.Tensor | None = None,
870
+ image_grid_thw: torch.Tensor | None = None,
871
+ video_grid_thw: torch.Tensor | None = None,
872
+ mm_token_type_ids: torch.Tensor | None = None,
873
+ audio_values: torch.Tensor | None = None,
874
+ video_values: torch.Tensor | None = None,
875
+ sensor_values: torch.Tensor | None = None,
876
+ prefix_mask: torch.Tensor | None = None,
877
+ reasoning_effort: str | None = None,
878
+ reasoning_budget: int | None = None,
879
+ use_cache: bool | None = None,
880
+ output_attentions: bool | None = None,
881
+ output_hidden_states: bool | None = None,
882
+ return_dict: bool | None = None,
883
+ collect_reasoning_diagnostics: bool = True,
884
+ cache_position: torch.Tensor | None = None,
885
+ logits_to_keep: int | torch.Tensor = 0,
886
+ vocab_indices: torch.Tensor | None = None,
887
+ **kwargs: Any,
888
+ ) -> DendroCausalLMOutput | tuple[Any, ...]:
889
+ return_dict = self.config.use_return_dict if return_dict is None else bool(return_dict)
890
+ base = self.model(
891
+ input_ids=input_ids,
892
+ attention_mask=attention_mask,
893
+ position_ids=position_ids,
894
+ past_key_values=past_key_values,
895
+ inputs_embeds=inputs_embeds,
896
+ pixel_values=pixel_values,
897
+ pixel_values_videos=pixel_values_videos,
898
+ image_grid_thw=image_grid_thw,
899
+ video_grid_thw=video_grid_thw,
900
+ mm_token_type_ids=mm_token_type_ids,
901
+ audio_values=audio_values,
902
+ video_values=video_values,
903
+ sensor_values=sensor_values,
904
+ prefix_mask=prefix_mask,
905
+ labels=labels,
906
+ reasoning_effort=reasoning_effort,
907
+ reasoning_budget=reasoning_budget,
908
+ use_cache=use_cache,
909
+ output_attentions=output_attentions,
910
+ output_hidden_states=output_hidden_states,
911
+ return_dict=True,
912
+ collect_reasoning_diagnostics=collect_reasoning_diagnostics,
913
+ cache_position=cache_position,
914
+ **kwargs,
915
+ )
916
+ assert isinstance(base, DendroBaseModelOutputWithPast)
917
+ hidden = base.last_hidden_state
918
+ assert hidden is not None
919
+ if isinstance(logits_to_keep, int) and logits_to_keep > 0:
920
+ hidden_for_logits = hidden[:, -logits_to_keep:]
921
+ elif torch.is_tensor(logits_to_keep):
922
+ hidden_for_logits = hidden.index_select(1, logits_to_keep.to(hidden.device))
923
+ else:
924
+ hidden_for_logits = hidden
925
+ if vocab_indices is None:
926
+ logits = self.source_layer.tied_logits(
927
+ hidden_for_logits, vocab_size=self.config.vocab_size
928
+ ).float()
929
+ else:
930
+ logits = self.source_layer.tied_selected_logits(
931
+ hidden_for_logits,
932
+ vocab_indices=vocab_indices,
933
+ vocab_size=self.config.vocab_size,
934
+ ).float()
935
+
936
+ loss = None
937
+ aligned_labels = base.aligned_labels
938
+ if labels is not None:
939
+ if vocab_indices is not None:
940
+ raise ValueError("vocab_indices is a trainer-only projection and cannot be combined with model labels")
941
+ if aligned_labels is None:
942
+ raise RuntimeError("labels were supplied but not aligned by the input projector")
943
+ full_logits = logits
944
+ if full_logits.shape[1] != aligned_labels.shape[1]:
945
+ full_logits = self.source_layer.tied_logits(hidden, vocab_size=self.config.vocab_size).float()
946
+ shift_logits = full_logits[:, :-1, :].contiguous()
947
+ shift_labels = aligned_labels[:, 1:].to(shift_logits.device).contiguous()
948
+ loss = F.cross_entropy(
949
+ shift_logits.view(-1, self.config.vocab_size),
950
+ shift_labels.view(-1),
951
+ ignore_index=-100,
952
+ )
953
+
954
+ result = DendroCausalLMOutput(
955
+ loss=loss,
956
+ logits=logits,
957
+ past_key_values=base.past_key_values,
958
+ hidden_states=base.hidden_states,
959
+ attentions=base.attentions,
960
+ reasoning_state=base.reasoning_state,
961
+ modality_layout=base.modality_layout,
962
+ dendro_states=base.dendro_states,
963
+ )
964
+ if return_dict:
965
+ return result
966
+ return result.to_tuple()
967
+
968
+ @torch.no_grad()
969
+ def answer_image(
970
+ self,
971
+ processor: Any,
972
+ image: Any,
973
+ prompt: str,
974
+ *,
975
+ reasoning_effort: str = "high",
976
+ max_new_tokens: int | None = None,
977
+ do_sample: bool = False,
978
+ temperature: float = 0.7,
979
+ ) -> Any:
980
+ """Answer about an image with effort-scaled detail views.
981
+
982
+ The caller reuses this loaded model and its matching lightweight
983
+ processor. High/max add deterministic local views but never load a
984
+ second language model or vision tower.
985
+ """
986
+
987
+ from .vision_routing import answer_image_with_effort
988
+
989
+ return answer_image_with_effort(
990
+ self,
991
+ processor,
992
+ image,
993
+ prompt,
994
+ reasoning_effort=reasoning_effort,
995
+ max_new_tokens=max_new_tokens,
996
+ do_sample=do_sample,
997
+ temperature=temperature,
998
+ )
999
+
1000
+ def prepare_inputs_for_generation(
1001
+ self,
1002
+ input_ids: torch.Tensor,
1003
+ past_key_values: DendroKVCache | DendroTransplantCache | DendroAwakeningCache | None = None,
1004
+ attention_mask: torch.Tensor | None = None,
1005
+ inputs_embeds: torch.Tensor | None = None,
1006
+ cache_position: torch.Tensor | None = None,
1007
+ **kwargs: Any,
1008
+ ) -> dict[str, Any]:
1009
+ populated = past_key_values is not None and past_key_values.get_seq_length() > 0
1010
+ if populated:
1011
+ seen = past_key_values.total_tokens_seen
1012
+ if input_ids.shape[1] > seen:
1013
+ start = min(seen, input_ids.shape[1] - 1)
1014
+ input_ids = input_ids[:, start:]
1015
+ if attention_mask is not None:
1016
+ attention_mask = attention_mask[:, start:]
1017
+ else:
1018
+ input_ids = input_ids[:, -1:]
1019
+ if attention_mask is not None:
1020
+ attention_mask = attention_mask[:, -1:]
1021
+ inputs_embeds = None
1022
+ for key in ("pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw", "mm_token_type_ids", "audio_values", "video_values", "sensor_values", "prefix_mask"):
1023
+ kwargs[key] = None
1024
+ model_inputs: dict[str, Any] = {
1025
+ "input_ids": input_ids,
1026
+ "attention_mask": attention_mask,
1027
+ "past_key_values": past_key_values,
1028
+ "use_cache": kwargs.pop("use_cache", True),
1029
+ "cache_position": cache_position,
1030
+ }
1031
+ if inputs_embeds is not None and not populated:
1032
+ model_inputs["inputs_embeds"] = inputs_embeds
1033
+ model_inputs["input_ids"] = None
1034
+ for key in (
1035
+ "pixel_values",
1036
+ "pixel_values_videos",
1037
+ "image_grid_thw",
1038
+ "video_grid_thw",
1039
+ "mm_token_type_ids",
1040
+ "audio_values",
1041
+ "video_values",
1042
+ "sensor_values",
1043
+ "prefix_mask",
1044
+ "reasoning_effort",
1045
+ "reasoning_budget",
1046
+ "logits_to_keep",
1047
+ ):
1048
+ if key in kwargs:
1049
+ model_inputs[key] = kwargs[key]
1050
+ return model_inputs
1051
+
1052
+ @staticmethod
1053
+ def _reorder_cache(
1054
+ past_key_values: DendroKVCache | DendroTransplantCache | DendroAwakeningCache,
1055
+ beam_idx: torch.Tensor,
1056
+ ) -> DendroKVCache | DendroTransplantCache | DendroAwakeningCache:
1057
+ return past_key_values.reorder_cache(beam_idx)
1058
+
1059
+ @torch.no_grad()
1060
+ def _local_generate(
1061
+ self,
1062
+ input_ids: torch.Tensor,
1063
+ *,
1064
+ # A soft practical allowance: generation still exits immediately on EOS,
1065
+ # while normal structured answers are not silently clipped at 20 tokens.
1066
+ max_new_tokens: int = 384,
1067
+ completion_reserve_tokens: int | None = None,
1068
+ min_new_tokens: int = 0,
1069
+ attention_mask: torch.Tensor | None = None,
1070
+ do_sample: bool = False,
1071
+ temperature: float = 1.0,
1072
+ top_k: int = 0,
1073
+ top_p: float = 1.0,
1074
+ repetition_penalty: float = 1.0,
1075
+ eos_token_id: int | list[int] | tuple[int, ...] | None = None,
1076
+ pad_token_id: int | None = None,
1077
+ reasoning_effort: str | None = None,
1078
+ reasoning_budget: int | None = None,
1079
+ generator: torch.Generator | None = None,
1080
+ use_cache: bool = True,
1081
+ token_callback: Any | None = None,
1082
+ **model_kwargs: Any,
1083
+ ) -> torch.Tensor:
1084
+ if not torch.is_tensor(input_ids):
1085
+ raise TypeError("input_ids must be a torch.Tensor")
1086
+ if input_ids.ndim != 2:
1087
+ raise ValueError("input_ids must have shape (batch, sequence)")
1088
+ if input_ids.shape[1] == 0:
1089
+ raise ValueError("input_ids must contain at least one token")
1090
+ if isinstance(max_new_tokens, bool) or not isinstance(max_new_tokens, int):
1091
+ raise TypeError("max_new_tokens must be an integer")
1092
+ if max_new_tokens < 0:
1093
+ raise ValueError("max_new_tokens cannot be negative")
1094
+ if completion_reserve_tokens is None:
1095
+ completion_reserve_tokens = int(
1096
+ getattr(self.config, "generation_completion_reserve_tokens", 0)
1097
+ )
1098
+ if (
1099
+ isinstance(completion_reserve_tokens, bool)
1100
+ or not isinstance(completion_reserve_tokens, int)
1101
+ ):
1102
+ raise TypeError("completion_reserve_tokens must be an integer")
1103
+ if completion_reserve_tokens < 0:
1104
+ raise ValueError("completion_reserve_tokens cannot be negative")
1105
+ if isinstance(min_new_tokens, bool) or not isinstance(min_new_tokens, int):
1106
+ raise TypeError("min_new_tokens must be an integer")
1107
+ if min_new_tokens < 0:
1108
+ raise ValueError("min_new_tokens cannot be negative")
1109
+ if min_new_tokens > max_new_tokens:
1110
+ raise ValueError("min_new_tokens cannot exceed max_new_tokens")
1111
+ if isinstance(top_k, bool) or not isinstance(top_k, int):
1112
+ raise TypeError("top_k must be an integer")
1113
+ if top_k < 0:
1114
+ raise ValueError("top_k cannot be negative")
1115
+ try:
1116
+ temperature = float(temperature)
1117
+ except (TypeError, ValueError) as error:
1118
+ raise TypeError("temperature must be a real number") from error
1119
+ if not math.isfinite(temperature) or temperature <= 0.0:
1120
+ raise ValueError("temperature must be finite and positive")
1121
+ try:
1122
+ top_p = float(top_p)
1123
+ except (TypeError, ValueError) as error:
1124
+ raise TypeError("top_p must be a real number") from error
1125
+ if not math.isfinite(top_p) or not 0.0 < top_p <= 1.0:
1126
+ raise ValueError("top_p must be finite and in (0, 1]")
1127
+ try:
1128
+ repetition_penalty = float(repetition_penalty)
1129
+ except (TypeError, ValueError) as error:
1130
+ raise TypeError("repetition_penalty must be a real number") from error
1131
+ if not math.isfinite(repetition_penalty) or repetition_penalty <= 0.0:
1132
+ raise ValueError("repetition_penalty must be finite and positive")
1133
+
1134
+ # These forward-only knobs cannot be honored by a tensor-returning local
1135
+ # generation loop. Consume them so callers that pass generic model
1136
+ # kwargs cannot collide with the invariants below.
1137
+ model_kwargs.pop("return_dict", None)
1138
+ model_kwargs.pop("logits_to_keep", None)
1139
+ supported_forward_kwargs = {
1140
+ "audio_values",
1141
+ "cache_position",
1142
+ "image_grid_thw",
1143
+ "mm_token_type_ids",
1144
+ "output_attentions",
1145
+ "output_hidden_states",
1146
+ "pixel_values",
1147
+ "pixel_values_videos",
1148
+ "position_ids",
1149
+ "prefix_mask",
1150
+ "sensor_values",
1151
+ "video_grid_thw",
1152
+ "video_values",
1153
+ }
1154
+ unsupported = sorted(set(model_kwargs).difference(supported_forward_kwargs))
1155
+ if unsupported:
1156
+ names = ", ".join(unsupported)
1157
+ raise TypeError(f"unsupported local generation argument(s): {names}")
1158
+
1159
+ def normalize_eos_ids(value: Any) -> tuple[int, ...]:
1160
+ if value is None:
1161
+ return ()
1162
+ values = value if isinstance(value, (list, tuple)) else (value,)
1163
+ normalized: list[int] = []
1164
+ for token_id in values:
1165
+ if isinstance(token_id, bool) or not isinstance(token_id, int):
1166
+ raise TypeError("eos_token_id must be an integer or a list of integers")
1167
+ if token_id not in normalized:
1168
+ normalized.append(token_id)
1169
+ return tuple(normalized)
1170
+
1171
+ configured_eos = self.config.eos_token_id if eos_token_id is None else eos_token_id
1172
+ eos_ids = normalize_eos_ids(configured_eos)
1173
+ configured_pad = self.config.pad_token_id if pad_token_id is None else pad_token_id
1174
+ if configured_pad is None:
1175
+ configured_pad = next(
1176
+ (token_id for token_id in eos_ids if 0 <= token_id < self.config.vocab_size),
1177
+ 0,
1178
+ )
1179
+ if isinstance(configured_pad, bool) or not isinstance(configured_pad, int):
1180
+ raise TypeError("pad_token_id must be an integer")
1181
+ if not 0 <= configured_pad < self.config.vocab_size:
1182
+ raise ValueError(
1183
+ f"pad_token_id must be in [0, {self.config.vocab_size}), got {configured_pad}"
1184
+ )
1185
+ pad = configured_pad
1186
+ valid_eos_ids = tuple(
1187
+ token_id for token_id in eos_ids if 0 <= token_id < self.config.vocab_size
1188
+ )
1189
+ eos_tensor = torch.tensor(eos_ids, device=input_ids.device, dtype=input_ids.dtype)
1190
+ valid_eos_tensor = torch.tensor(
1191
+ valid_eos_ids, device=input_ids.device, dtype=torch.long
1192
+ )
1193
+
1194
+ if attention_mask is None:
1195
+ prompt_mask = torch.ones_like(input_ids, dtype=torch.bool)
1196
+ else:
1197
+ if not torch.is_tensor(attention_mask):
1198
+ attention_mask = torch.as_tensor(attention_mask, device=input_ids.device)
1199
+ else:
1200
+ attention_mask = attention_mask.to(device=input_ids.device)
1201
+ if attention_mask.ndim == 1 and input_ids.shape[0] == 1:
1202
+ attention_mask = attention_mask.unsqueeze(0)
1203
+ if attention_mask.shape != input_ids.shape:
1204
+ raise ValueError(
1205
+ "attention_mask must have the same (batch, sequence) shape as input_ids"
1206
+ )
1207
+ prompt_mask = attention_mask.to(dtype=torch.bool)
1208
+
1209
+ batch_size, prompt_length = input_ids.shape
1210
+ total_token_budget = max_new_tokens + completion_reserve_tokens
1211
+ capacity = prompt_length + total_token_budget
1212
+ token_buffer = input_ids.new_empty((batch_size, capacity))
1213
+ mask_buffer = torch.empty(
1214
+ (batch_size, capacity), device=input_ids.device, dtype=torch.bool
1215
+ )
1216
+ token_buffer[:, :prompt_length].copy_(input_ids)
1217
+ mask_buffer[:, :prompt_length].copy_(prompt_mask)
1218
+ generated_length = prompt_length
1219
+ generated = token_buffer[:, :generated_length]
1220
+ generated_mask = mask_buffer[:, :generated_length]
1221
+ current = input_ids
1222
+ current_mask = generated_mask
1223
+ cache: DendroKVCache | DendroTransplantCache | DendroAwakeningCache | None = None
1224
+ # ``use_cache`` is a public generation argument, not a model-forward
1225
+ # kwarg. Keeping it explicit prevents it from also arriving through
1226
+ # ``model_kwargs`` and colliding with the value passed below. The
1227
+ # uncached path deliberately recomputes the complete prefix each step.
1228
+ recompute_full_prefix = not use_cache
1229
+ finished = torch.zeros(input_ids.shape[0], device=input_ids.device, dtype=torch.bool)
1230
+ if hasattr(self.model, "transplant_core") and self.model.transplant_core is not None and hasattr(self.model.transplant_core, "prebind_weights"):
1231
+ self.model.transplant_core.prebind_weights(device=input_ids.device)
1232
+ prefix_kwargs = dict(model_kwargs)
1233
+ for _step in range(total_token_budget):
1234
+ if recompute_full_prefix:
1235
+ current = generated
1236
+ current_mask = generated_mask
1237
+ cache = None
1238
+ output = self(
1239
+ input_ids=current,
1240
+ attention_mask=current_mask,
1241
+ past_key_values=cache,
1242
+ reasoning_effort=reasoning_effort,
1243
+ reasoning_budget=reasoning_budget,
1244
+ use_cache=not recompute_full_prefix,
1245
+ return_dict=True,
1246
+ logits_to_keep=1,
1247
+ collect_reasoning_diagnostics=False,
1248
+ **prefix_kwargs,
1249
+ )
1250
+ assert isinstance(output, DendroCausalLMOutput)
1251
+ cache = None if recompute_full_prefix else output.past_key_values
1252
+ logits = output.logits[:, -1, :]
1253
+ if repetition_penalty != 1.0:
1254
+ previous_scores = logits.gather(1, generated)
1255
+ penalized_scores = torch.where(
1256
+ previous_scores < 0,
1257
+ previous_scores * repetition_penalty,
1258
+ previous_scores / repetition_penalty,
1259
+ )
1260
+ logits = logits.scatter(1, generated, penalized_scores)
1261
+ if _step < min_new_tokens and valid_eos_tensor.numel() > 0:
1262
+ logits = logits.index_fill(1, valid_eos_tensor, -torch.inf)
1263
+ if do_sample:
1264
+ logits = logits / temperature
1265
+ if top_k > 0:
1266
+ candidate_logits, candidate_indices = logits.topk(
1267
+ min(top_k, logits.shape[-1]), dim=-1
1268
+ )
1269
+ if top_p < 1.0:
1270
+ cumulative = torch.softmax(candidate_logits, dim=-1).cumsum(dim=-1)
1271
+ remove = cumulative > top_p
1272
+ remove[..., 1:] = remove[..., :-1].clone()
1273
+ remove[..., 0] = False
1274
+ candidate_logits = candidate_logits.masked_fill(remove, -torch.inf)
1275
+ candidate_position = torch.multinomial(
1276
+ torch.softmax(candidate_logits, dim=-1), 1, generator=generator
1277
+ )
1278
+ next_token = candidate_indices.gather(-1, candidate_position)
1279
+ elif top_p < 1.0:
1280
+ sorted_logits, sorted_indices = logits.sort(descending=True)
1281
+ cumulative = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
1282
+ remove = cumulative > top_p
1283
+ remove[..., 1:] = remove[..., :-1].clone()
1284
+ remove[..., 0] = False
1285
+ sorted_logits = sorted_logits.masked_fill(remove, -torch.inf)
1286
+ sorted_position = torch.multinomial(
1287
+ torch.softmax(sorted_logits, dim=-1), 1, generator=generator
1288
+ )
1289
+ next_token = sorted_indices.gather(-1, sorted_position)
1290
+ else:
1291
+ next_token = torch.multinomial(
1292
+ torch.softmax(logits, dim=-1), 1, generator=generator
1293
+ )
1294
+ else:
1295
+ next_token = logits.argmax(dim=-1, keepdim=True)
1296
+ next_token = torch.where(finished.unsqueeze(-1), torch.full_like(next_token, pad), next_token)
1297
+ active = ~finished
1298
+ if token_callback is not None:
1299
+ token_callback(next_token.detach(), _step + 1, active.detach())
1300
+ token_buffer[:, generated_length].copy_(next_token.squeeze(-1))
1301
+ mask_buffer[:, generated_length].copy_(active)
1302
+ generated_length += 1
1303
+ generated = token_buffer[:, :generated_length]
1304
+ generated_mask = mask_buffer[:, :generated_length]
1305
+ if eos_tensor.numel() > 0:
1306
+ hit_eos = (
1307
+ next_token.squeeze(-1).unsqueeze(-1) == eos_tensor.unsqueeze(0)
1308
+ ).any(dim=-1)
1309
+ finished = finished | hit_eos
1310
+ if bool(finished.all().item()):
1311
+ break
1312
+ if not recompute_full_prefix:
1313
+ current = next_token
1314
+ # Rows that finished on an earlier iteration must not add their
1315
+ # padding tokens to the cache as valid context. ``active`` was
1316
+ # captured before this step's EOS update, so a newly emitted EOS
1317
+ # is still committed once while older finished rows stay masked.
1318
+ current_mask = active.unsqueeze(-1)
1319
+ prefix_kwargs = {}
1320
+ return token_buffer[:, :generated_length].contiguous()
1321
+
1322
+ def _resolve_thinking_token_ids(
1323
+ self,
1324
+ start_token_id: int | None,
1325
+ end_token_id: int | None,
1326
+ ) -> tuple[int, int]:
1327
+ """Resolve configured thinking boundaries or the transplanted Qwen IDs."""
1328
+
1329
+ start = self.config.thinking_start_token_id if start_token_id is None else start_token_id
1330
+ end = self.config.thinking_end_token_id if end_token_id is None else end_token_id
1331
+ # Qwen3.5's copied tokenizer stores these as ordinary added tokens rather
1332
+ # than ``special`` tokens, so skip_special_tokens cannot be used as a
1333
+ # privacy boundary. They are inferred only when the vocabulary covers the
1334
+ # known IDs; smaller/native vocabularies must configure boundaries.
1335
+ if start is None and self.config.vocab_size > 248_068:
1336
+ start = 248_068
1337
+ if end is None and self.config.vocab_size > 248_069:
1338
+ end = 248_069
1339
+ if start is None or end is None:
1340
+ raise ValueError(
1341
+ "deliberative generation requires thinking_start_token_id and "
1342
+ "thinking_end_token_id for this vocabulary"
1343
+ )
1344
+ for name, value in (("thinking_start_token_id", start), ("thinking_end_token_id", end)):
1345
+ if isinstance(value, bool) or not isinstance(value, int):
1346
+ raise TypeError(f"{name} must be an integer")
1347
+ if not 0 <= value < self.config.vocab_size:
1348
+ raise ValueError(f"{name} must be in [0, {self.config.vocab_size})")
1349
+ if start == end:
1350
+ raise ValueError("thinking start and end token IDs must be distinct")
1351
+ return int(start), int(end)
1352
+
1353
+ @torch.no_grad()
1354
+ def generate_answer(
1355
+ self,
1356
+ input_ids: torch.Tensor,
1357
+ *,
1358
+ attention_mask: torch.Tensor | None = None,
1359
+ max_reasoning_tokens: int | None = None,
1360
+ min_reasoning_tokens: int | None = None,
1361
+ max_answer_tokens: int | None = None,
1362
+ reasoning_effort: str | None = None,
1363
+ reasoning_budget: int | None = None,
1364
+ thinking_start_token_id: int | None = None,
1365
+ thinking_end_token_id: int | None = None,
1366
+ do_sample: bool = False,
1367
+ temperature: float = 1.0,
1368
+ top_k: int = 0,
1369
+ top_p: float = 1.0,
1370
+ repetition_penalty: float = 1.0,
1371
+ answer_eos_token_id: int | list[int] | tuple[int, ...] | None = None,
1372
+ pad_token_id: int | None = None,
1373
+ generator: torch.Generator | None = None,
1374
+ use_cache: bool = True,
1375
+ progress_callback: Any | None = None,
1376
+ answer_token_callback: Any | None = None,
1377
+ **model_kwargs: Any,
1378
+ ) -> DendroAnswerGenerationOutput:
1379
+ """Deliberate privately, then return only independently budgeted answers.
1380
+
1381
+ The first phase ends at the configured ``</think>`` token. An explicit
1382
+ ``max_reasoning_tokens`` remains a hard caller limit. With the packaged
1383
+ ``adaptive_context`` policy and no explicit maximum, private reasoning may
1384
+ expand past the effort profile up to the real context boundary while the
1385
+ complete visible-answer allocation stays reserved. If that physical
1386
+ boundary is reached, ``</think>`` is injected before answer generation.
1387
+
1388
+ Rows are processed independently. This avoids leaking one batch member's
1389
+ early-close padding into another member's cache and permits variable hidden
1390
+ reasoning lengths without exposing them in a rectangular result tensor.
1391
+ ``progress_callback`` receives phase/count/budget metadata only; private
1392
+ token IDs are never exposed. ``answer_token_callback`` receives visible
1393
+ answer token IDs as they are generated, enabling privacy-safe streaming.
1394
+ """
1395
+
1396
+ if not torch.is_tensor(input_ids) or input_ids.ndim != 2:
1397
+ raise ValueError("input_ids must be a tensor with shape (batch, sequence)")
1398
+ if input_ids.shape[0] == 0 or input_ids.shape[1] == 0:
1399
+ raise ValueError("input_ids must contain at least one token per row")
1400
+ for name, value in (("max_reasoning_tokens", max_reasoning_tokens),
1401
+ ("min_reasoning_tokens", min_reasoning_tokens)):
1402
+ if value is None:
1403
+ continue
1404
+ if isinstance(value, bool) or not isinstance(value, int):
1405
+ raise TypeError(f"{name} must be an integer")
1406
+ if value < 0:
1407
+ raise ValueError(f"{name} cannot be negative")
1408
+ explicit_answer_max = max_answer_tokens is not None
1409
+ if max_answer_tokens is not None:
1410
+ if isinstance(max_answer_tokens, bool) or not isinstance(max_answer_tokens, int):
1411
+ raise TypeError("max_answer_tokens must be an integer")
1412
+ if max_answer_tokens < 0:
1413
+ raise ValueError("max_answer_tokens cannot be negative")
1414
+ configured_answer_reserve = int(
1415
+ getattr(self.config, "generation_default_max_answer_tokens", 384)
1416
+ )
1417
+ if max_answer_tokens is None:
1418
+ max_answer_tokens = configured_answer_reserve
1419
+ # Validate even when both generation phases have a zero token allowance.
1420
+ self.config.resolve_reasoning_budget(reasoning_effort, reasoning_budget)
1421
+ resolved_reasoning_min, resolved_reasoning_max = self.config.reasoning_token_budget(
1422
+ reasoning_effort,
1423
+ maximum=max_reasoning_tokens,
1424
+ minimum=min_reasoning_tokens,
1425
+ )
1426
+ answer_reasoning_budget = self.config.answer_commit_budget(
1427
+ reasoning_effort, reasoning_budget
1428
+ )
1429
+ thinking_start, thinking_end = self._resolve_thinking_token_ids(
1430
+ thinking_start_token_id, thinking_end_token_id
1431
+ )
1432
+ if any(key in model_kwargs for key in ("position_ids", "cache_position")):
1433
+ raise TypeError(
1434
+ "generate_answer derives private-phase positions and does not accept "
1435
+ "position_ids or cache_position"
1436
+ )
1437
+
1438
+ if attention_mask is None:
1439
+ supplied_mask = torch.ones_like(input_ids, dtype=torch.bool)
1440
+ else:
1441
+ if not torch.is_tensor(attention_mask):
1442
+ attention_mask = torch.as_tensor(attention_mask, device=input_ids.device)
1443
+ supplied_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
1444
+ if supplied_mask.shape != input_ids.shape:
1445
+ raise ValueError("attention_mask must have the same shape as input_ids")
1446
+
1447
+ configured_eos = (
1448
+ self.config.eos_token_id
1449
+ if answer_eos_token_id is None
1450
+ else answer_eos_token_id
1451
+ )
1452
+ eos_values = configured_eos if isinstance(configured_eos, (list, tuple)) else (configured_eos,)
1453
+ normalized_eos: tuple[int, ...] = tuple(
1454
+ int(value)
1455
+ for value in eos_values
1456
+ if value is not None and not isinstance(value, bool) and isinstance(value, int)
1457
+ )
1458
+ resolved_pad = self.config.pad_token_id if pad_token_id is None else pad_token_id
1459
+ if resolved_pad is None:
1460
+ resolved_pad = next(
1461
+ (value for value in normalized_eos if 0 <= value < self.config.vocab_size),
1462
+ 0,
1463
+ )
1464
+ if isinstance(resolved_pad, bool) or not isinstance(resolved_pad, int):
1465
+ raise TypeError("pad_token_id must be an integer")
1466
+ if not 0 <= resolved_pad < self.config.vocab_size:
1467
+ raise ValueError("pad_token_id must be in [0, vocab_size)")
1468
+
1469
+ if input_ids.shape[0] > 1 and any(
1470
+ model_kwargs.get(name) is not None
1471
+ for name in (
1472
+ "pixel_values",
1473
+ "pixel_values_videos",
1474
+ "audio_values",
1475
+ "video_values",
1476
+ "sensor_values",
1477
+ )
1478
+ ):
1479
+ raise ValueError(
1480
+ "batched deliberative multimodal generation is not yet supported; "
1481
+ "generate each multimodal request separately"
1482
+ )
1483
+
1484
+ answer_rows: list[torch.Tensor] = []
1485
+ reasoning_counts: list[int] = []
1486
+ reasoning_minimums: list[int] = []
1487
+ reasoning_budgets: list[int] = []
1488
+ reasoning_completed: list[bool] = []
1489
+ forced_close: list[bool] = []
1490
+ answer_counts: list[int] = []
1491
+ answer_eos_reached: list[bool] = []
1492
+ answer_reasoning_filtered: list[bool] = []
1493
+ finish_reasons: list[str] = []
1494
+ direct_final = self.config.reasoning_effort_profile(reasoning_effort)[0] == "direct"
1495
+
1496
+ batch_size = input_ids.shape[0]
1497
+ for name, callback in (
1498
+ ("progress_callback", progress_callback),
1499
+ ("answer_token_callback", answer_token_callback),
1500
+ ):
1501
+ if callback is not None and not callable(callback):
1502
+ raise TypeError(f"{name} must be callable")
1503
+ sampling_kwargs = {
1504
+ "do_sample": do_sample,
1505
+ "temperature": temperature,
1506
+ "top_k": top_k,
1507
+ "top_p": top_p,
1508
+ "repetition_penalty": repetition_penalty,
1509
+ "generator": generator,
1510
+ "use_cache": use_cache,
1511
+ }
1512
+ for row_index in range(batch_size):
1513
+ valid = supplied_mask[row_index]
1514
+ if not bool(valid.any().item()):
1515
+ raise ValueError("each input row must contain at least one unmasked token")
1516
+ row_prompt = input_ids[row_index, valid].unsqueeze(0)
1517
+ row_mask = torch.ones_like(row_prompt, dtype=torch.bool)
1518
+ row_kwargs: dict[str, Any] = {}
1519
+ for key, value in model_kwargs.items():
1520
+ if torch.is_tensor(value) and value.ndim > 0 and value.shape[0] == batch_size:
1521
+ sliced = value[row_index : row_index + 1]
1522
+ if (
1523
+ key in {"mm_token_type_ids", "prefix_mask"}
1524
+ and sliced.ndim == 2
1525
+ and sliced.shape[1] == input_ids.shape[1]
1526
+ ):
1527
+ sliced = sliced[:, valid.to(device=sliced.device)]
1528
+ row_kwargs[key] = sliced
1529
+ else:
1530
+ row_kwargs[key] = value
1531
+
1532
+ def sequence_aligned_kwargs(length: int) -> dict[str, Any]:
1533
+ aligned = dict(row_kwargs)
1534
+ for key in ("mm_token_type_ids", "prefix_mask"):
1535
+ value = aligned.get(key)
1536
+ if value is None or not torch.is_tensor(value) or value.ndim != 2:
1537
+ continue
1538
+ if value.shape[1] > length:
1539
+ raise ValueError(f"{key} is longer than the deliberative sequence")
1540
+ if value.shape[1] < length:
1541
+ padding = value.new_zeros((value.shape[0], length - value.shape[1]))
1542
+ aligned[key] = torch.cat((value, padding), dim=1)
1543
+ return aligned
1544
+
1545
+ prompt_values = row_prompt[0].detach().cpu().tolist()
1546
+ last_start = max(
1547
+ (index for index, token in enumerate(prompt_values) if token == thinking_start),
1548
+ default=-1,
1549
+ )
1550
+ last_end = max(
1551
+ (index for index, token in enumerate(prompt_values) if token == thinking_end),
1552
+ default=-1,
1553
+ )
1554
+ row_reasoning_max = resolved_reasoning_max
1555
+ if (
1556
+ not direct_final
1557
+ and
1558
+ max_reasoning_tokens is None
1559
+ and self.config.reasoning_token_budget_policy == "adaptive_context"
1560
+ and self.config.private_deliberation_trained
1561
+ ):
1562
+ context_limits = [int(self.config.max_position_embeddings)]
1563
+ if self.config.transplant_mode in {"exact", "awakening"}:
1564
+ context_limits.append(int(self.config.transplant_attention_window))
1565
+ if self.config.max_cache_length is not None:
1566
+ context_limits.append(int(self.config.max_cache_length))
1567
+ if self.config.sliding_window is not None:
1568
+ context_limits.append(int(self.config.sliding_window))
1569
+ # Reasoning past the active attention/cache window is not deeper:
1570
+ # early private tokens have already been evicted. The minimum
1571
+ # physical limit is therefore the honest adaptive boundary.
1572
+ active_context_capacity = min(context_limits)
1573
+ physical_allowance = (
1574
+ active_context_capacity
1575
+ - int(row_prompt.shape[1])
1576
+ - int(max_answer_tokens)
1577
+ - int(
1578
+ getattr(
1579
+ self.config,
1580
+ "generation_completion_reserve_tokens",
1581
+ 0,
1582
+ )
1583
+ )
1584
+ - int(self.config.reasoning_context_guard_tokens)
1585
+ - 1
1586
+ )
1587
+ if physical_allowance < resolved_reasoning_min:
1588
+ raise ValueError(
1589
+ "prompt plus reserved answer leaves insufficient active context for "
1590
+ "the configured minimum private reasoning budget"
1591
+ )
1592
+ # The active window is authoritative. Effort profiles are the
1593
+ # initial/minimum policy, never permission to consume the answer
1594
+ # allocation or silently evict the user's goal before answering.
1595
+ row_reasoning_max = physical_allowance
1596
+ if row_reasoning_max > 0 and last_start <= last_end:
1597
+ opener = row_prompt.new_tensor([[thinking_start]])
1598
+ row_prompt = torch.cat((row_prompt, opener), dim=1)
1599
+ row_mask = torch.cat((row_mask, torch.ones_like(opener, dtype=torch.bool)), dim=1)
1600
+
1601
+ reasoning_prefix_length = row_prompt.shape[1]
1602
+ if progress_callback is not None:
1603
+ progress_callback("reasoning_start", 0, row_reasoning_max)
1604
+ if row_reasoning_max:
1605
+ def private_progress(
1606
+ _tokens: torch.Tensor,
1607
+ count: int,
1608
+ _active: torch.Tensor,
1609
+ ) -> None:
1610
+ if progress_callback is not None:
1611
+ progress_callback("reasoning", count, row_reasoning_max)
1612
+
1613
+ deliberated = self._local_generate(
1614
+ row_prompt,
1615
+ attention_mask=row_mask,
1616
+ max_new_tokens=row_reasoning_max,
1617
+ completion_reserve_tokens=0,
1618
+ min_new_tokens=resolved_reasoning_min,
1619
+ eos_token_id=thinking_end,
1620
+ pad_token_id=resolved_pad,
1621
+ reasoning_effort=reasoning_effort,
1622
+ reasoning_budget=reasoning_budget,
1623
+ token_callback=private_progress,
1624
+ **sampling_kwargs,
1625
+ **sequence_aligned_kwargs(row_prompt.shape[1]),
1626
+ )
1627
+ else:
1628
+ deliberated = row_prompt
1629
+ reasoning_suffix = deliberated[0, reasoning_prefix_length:]
1630
+ close_positions = torch.nonzero(
1631
+ reasoning_suffix.eq(thinking_end), as_tuple=False
1632
+ ).flatten()
1633
+ closed_naturally = close_positions.numel() > 0
1634
+ if row_reasoning_max == 0:
1635
+ # Low/direct effort does not manufacture a private phase. Close
1636
+ # only a caller-supplied dangling think span, then preserve the
1637
+ # entire independently budgeted answer allowance.
1638
+ if last_start > last_end:
1639
+ closer = deliberated.new_tensor([[thinking_end]])
1640
+ internal = torch.cat((deliberated, closer), dim=1)
1641
+ else:
1642
+ internal = deliberated
1643
+ closed_naturally = True
1644
+ reasoning_count = 0
1645
+ elif closed_naturally:
1646
+ close_offset = int(close_positions[0].item())
1647
+ reasoning_count = close_offset
1648
+ internal = deliberated[:, : reasoning_prefix_length + close_offset + 1]
1649
+ else:
1650
+ reasoning_count = int(reasoning_suffix.numel())
1651
+ closer = deliberated.new_tensor([[thinking_end]])
1652
+ internal = torch.cat((deliberated, closer), dim=1)
1653
+
1654
+ internal_mask = torch.ones_like(internal, dtype=torch.bool)
1655
+ row_answer_max = max_answer_tokens
1656
+ adaptive_answer = (
1657
+ not explicit_answer_max
1658
+ and getattr(self.config, "answer_token_budget_policy", "fixed")
1659
+ == "adaptive_context"
1660
+ )
1661
+ if adaptive_answer:
1662
+ # Logical request/output capacity is distinct from the bounded
1663
+ # physical KV attention window. Sliding caches cap memory while
1664
+ # absolute positions and total generation may continue. Private
1665
+ # reasoning remains physically guarded above so its answer-phase
1666
+ # prefill cannot grow without bound.
1667
+ active_capacity = min(
1668
+ int(self.config.max_position_embeddings),
1669
+ int(
1670
+ getattr(
1671
+ self.config,
1672
+ "generation_logical_context_tokens",
1673
+ self.config.max_position_embeddings,
1674
+ )
1675
+ ),
1676
+ )
1677
+ row_answer_max = max(
1678
+ 0,
1679
+ active_capacity
1680
+ - int(internal.shape[1])
1681
+ - int(self.config.reasoning_context_guard_tokens),
1682
+ )
1683
+ if progress_callback is not None:
1684
+ progress_callback("answer_start", 0, row_answer_max)
1685
+ if row_answer_max:
1686
+ streaming_inside_reasoning = False
1687
+
1688
+ def visible_progress(
1689
+ tokens: torch.Tensor,
1690
+ count: int,
1691
+ active: torch.Tensor,
1692
+ ) -> None:
1693
+ nonlocal streaming_inside_reasoning
1694
+ if progress_callback is not None:
1695
+ progress_callback("answer", count, row_answer_max)
1696
+ token = int(tokens[0, 0].item())
1697
+ if token == thinking_start:
1698
+ streaming_inside_reasoning = True
1699
+ elif token == thinking_end:
1700
+ streaming_inside_reasoning = False
1701
+ elif answer_token_callback is not None and not streaming_inside_reasoning:
1702
+ answer_token_callback(tokens, active)
1703
+
1704
+ answered = self._local_generate(
1705
+ internal,
1706
+ attention_mask=internal_mask,
1707
+ max_new_tokens=row_answer_max,
1708
+ completion_reserve_tokens=(
1709
+ 0
1710
+ if adaptive_answer or explicit_answer_max
1711
+ else int(
1712
+ getattr(
1713
+ self.config,
1714
+ "generation_completion_reserve_tokens",
1715
+ 0,
1716
+ )
1717
+ )
1718
+ ),
1719
+ eos_token_id=configured_eos,
1720
+ pad_token_id=resolved_pad,
1721
+ reasoning_effort=reasoning_effort,
1722
+ reasoning_budget=answer_reasoning_budget,
1723
+ token_callback=visible_progress,
1724
+ **sampling_kwargs,
1725
+ **sequence_aligned_kwargs(internal.shape[1]),
1726
+ )
1727
+ raw_answer = answered[0, internal.shape[1]:].contiguous()
1728
+ else:
1729
+ raw_answer = internal.new_empty((0,), dtype=internal.dtype)
1730
+ reached_eos = any(
1731
+ bool(raw_answer.eq(value).any().item()) for value in normalized_eos
1732
+ )
1733
+ # The thinking delimiters are ordinary Qwen added tokens, not special
1734
+ # tokens. Filter any accidentally reopened think span before exposing
1735
+ # answer IDs; an unclosed span is dropped through the end of the row.
1736
+ visible_tokens: list[int] = []
1737
+ inside_reasoning = False
1738
+ filtered_reasoning = False
1739
+ for token in raw_answer.detach().cpu().tolist():
1740
+ if token == thinking_start:
1741
+ inside_reasoning = True
1742
+ filtered_reasoning = True
1743
+ continue
1744
+ if token == thinking_end:
1745
+ inside_reasoning = False
1746
+ filtered_reasoning = True
1747
+ continue
1748
+ if not inside_reasoning:
1749
+ visible_tokens.append(int(token))
1750
+ answer = raw_answer.new_tensor(visible_tokens)
1751
+
1752
+ answer_rows.append(answer)
1753
+ reasoning_counts.append(reasoning_count)
1754
+ reasoning_minimums.append(resolved_reasoning_min)
1755
+ reasoning_budgets.append(row_reasoning_max)
1756
+ reasoning_completed.append(closed_naturally)
1757
+ forced_close.append(not closed_naturally)
1758
+ answer_counts.append(int(answer.numel()))
1759
+ answer_eos_reached.append(reached_eos)
1760
+ answer_reasoning_filtered.append(filtered_reasoning)
1761
+ finish_reasons.append("eos" if reached_eos else "length")
1762
+ if progress_callback is not None:
1763
+ progress_callback("complete", int(answer.numel()), row_answer_max)
1764
+
1765
+ maximum_answer_length = max(answer_counts, default=0)
1766
+ answer_token_ids = input_ids.new_full(
1767
+ (batch_size, maximum_answer_length), resolved_pad
1768
+ )
1769
+ answer_attention_mask = torch.zeros(
1770
+ (batch_size, maximum_answer_length),
1771
+ device=input_ids.device,
1772
+ dtype=torch.bool,
1773
+ )
1774
+ for row_index, answer in enumerate(answer_rows):
1775
+ count = answer.numel()
1776
+ if count:
1777
+ answer_token_ids[row_index, :count].copy_(answer)
1778
+ answer_attention_mask[row_index, :count] = True
1779
+
1780
+ return DendroAnswerGenerationOutput(
1781
+ answer_token_ids=answer_token_ids,
1782
+ answer_attention_mask=answer_attention_mask,
1783
+ reasoning_token_counts=tuple(reasoning_counts),
1784
+ reasoning_minimum_tokens=tuple(reasoning_minimums),
1785
+ reasoning_token_budgets=tuple(reasoning_budgets),
1786
+ reasoning_completed=tuple(reasoning_completed),
1787
+ reasoning_forced_close=tuple(forced_close),
1788
+ answer_token_counts=tuple(answer_counts),
1789
+ answer_eos_reached=tuple(answer_eos_reached),
1790
+ answer_reasoning_filtered=tuple(answer_reasoning_filtered),
1791
+ finish_reasons=tuple(finish_reasons),
1792
+ )
1793
+
1794
+ def generate(self, *args: Any, **kwargs: Any) -> torch.Tensor:
1795
+ private_reasoning = kwargs.pop("private_reasoning", None)
1796
+ if private_reasoning is None:
1797
+ private_reasoning = bool(
1798
+ getattr(self.config, "default_private_reasoning_generation", False)
1799
+ and self.config.private_deliberation_trained
1800
+ )
1801
+ if not isinstance(private_reasoning, bool):
1802
+ raise TypeError("private_reasoning must be a boolean")
1803
+ if private_reasoning:
1804
+ if args:
1805
+ input_ids, *remaining = args
1806
+ if remaining:
1807
+ raise TypeError(
1808
+ "Private generate accepts only input_ids as a positional argument"
1809
+ )
1810
+ else:
1811
+ input_ids = kwargs.pop("input_ids")
1812
+ if "max_new_tokens" in kwargs:
1813
+ max_answer_tokens = kwargs.pop("max_new_tokens")
1814
+ elif getattr(self.config, "answer_token_budget_policy", "fixed") == "adaptive_context":
1815
+ max_answer_tokens = None
1816
+ else:
1817
+ max_answer_tokens = int(
1818
+ getattr(self.config, "generation_default_max_answer_tokens", 384)
1819
+ )
1820
+ min_new_tokens = kwargs.pop("min_new_tokens", 0)
1821
+ if min_new_tokens:
1822
+ raise TypeError(
1823
+ "min_new_tokens is not supported by private generation; use "
1824
+ "generate_answer for explicit phase control"
1825
+ )
1826
+ answer_eos = kwargs.pop("eos_token_id", None)
1827
+ answer = self.generate_answer(
1828
+ input_ids,
1829
+ max_answer_tokens=max_answer_tokens,
1830
+ answer_eos_token_id=answer_eos,
1831
+ **kwargs,
1832
+ )
1833
+ return torch.cat(
1834
+ (input_ids, answer.answer_token_ids.to(device=input_ids.device)), dim=1
1835
+ ).contiguous()
1836
+ # The transplant path owns an exact hybrid cache that generic Transformers
1837
+ # cannot construct. Use the local loop so donor full-attention, gated-delta,
1838
+ # convolution, and Dendro recurrent states advance one token at a time.
1839
+ if self.model.transplant_core is not None and self.config.transplant_mode in {"exact", "awakening"}:
1840
+ if args:
1841
+ input_ids, *remaining = args
1842
+ if remaining:
1843
+ raise TypeError("Transplant generate accepts only input_ids as a positional argument")
1844
+ else:
1845
+ input_ids = kwargs.pop("input_ids")
1846
+ return self._local_generate(input_ids, **kwargs)
1847
+ if TRANSFORMERS_AVAILABLE:
1848
+ return GenerationMixin.generate(self, *args, **kwargs)
1849
+ if args:
1850
+ input_ids, *remaining = args
1851
+ if remaining:
1852
+ raise TypeError("Local generate accepts only input_ids as a positional argument")
1853
+ else:
1854
+ input_ids = kwargs.pop("input_ids")
1855
+ return self._local_generate(input_ids, **kwargs)
1856
+
1857
+ def architecture_audit(self, *, raise_on_error: bool = False) -> dict[str, Any]:
1858
+ parameters = list(self.named_parameters())
1859
+ private_layers = [
1860
+ name
1861
+ for name, module in self.named_modules()
1862
+ if isinstance(module, (nn.Linear, nn.Embedding))
1863
+ ]
1864
+ cell_count = sum(isinstance(module, DendroRecurrentCell) for module in self.modules())
1865
+ result = {
1866
+ **self.source_layer.audit(),
1867
+ "model_parameter_tensors": len(parameters),
1868
+ "model_parameter_elements": sum(parameter.numel() for _name, parameter in parameters),
1869
+ "parameter_names": [name for name, _parameter in parameters],
1870
+ "private_linear_or_embedding_modules": private_layers,
1871
+ "physical_recurrent_cells": cell_count,
1872
+ "virtual_base_depth": self.config.num_hidden_layers,
1873
+ "shared_input_output_token_view": True,
1874
+ "passes": len(parameters) == 1 and not private_layers and cell_count == 1,
1875
+ }
1876
+ if raise_on_error and not result["passes"]:
1877
+ raise RuntimeError(f"Dendro single-source audit failed: {result}")
1878
+ return result
1879
+
1880
+ def compile_for_inference(
1881
+ self,
1882
+ *,
1883
+ mode: str = "reduce-overhead",
1884
+ fullgraph: bool = False,
1885
+ dynamic: bool = True,
1886
+ ) -> "DendroForCausalLM":
1887
+ """Compile the recurrent cell with ``torch.compile`` when available."""
1888
+
1889
+ if not hasattr(torch, "compile"):
1890
+ raise RuntimeError("This PyTorch build does not provide torch.compile")
1891
+ self.model.cell = torch.compile( # type: ignore[assignment]
1892
+ self.model.cell,
1893
+ mode=mode,
1894
+ fullgraph=fullgraph,
1895
+ dynamic=dynamic,
1896
+ )
1897
+ return self
1898
+
1899
+
1900
+ def register_dendro_auto_classes() -> bool:
1901
+ """Register the model with local Hugging Face Auto classes."""
1902
+
1903
+ if not TRANSFORMERS_AVAILABLE:
1904
+ return False
1905
+ from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
1906
+
1907
+ from .tokenization_dendro_omni import DendroByteTokenizer
1908
+
1909
+ try:
1910
+ AutoConfig.register(DendroOmniConfig.model_type, DendroOmniConfig)
1911
+ except ValueError:
1912
+ pass
1913
+ try:
1914
+ AutoModel.register(DendroOmniConfig, DendroOmniModel)
1915
+ except ValueError:
1916
+ pass
1917
+ try:
1918
+ AutoModelForCausalLM.register(DendroOmniConfig, DendroForCausalLM)
1919
+ except ValueError:
1920
+ pass
1921
+ try:
1922
+ AutoTokenizer.register(
1923
+ DendroOmniConfig,
1924
+ slow_tokenizer_class=DendroByteTokenizer,
1925
+ fast_tokenizer_class=None,
1926
+ )
1927
+ except ValueError:
1928
+ pass
1929
+ return True
preprocessor_config.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "processor_class": "DendroVisionProcessor",
3
+ "auto_map": {
4
+ "AutoProcessor": "processing_dendro_omni.DendroVisionProcessor"
5
+ },
6
+ "image_processor_type": "DendroVisionProcessor",
7
+ "patch_size": 16,
8
+ "temporal_patch_size": 2,
9
+ "merge_size": 2,
10
+ "image_mean": [0.5, 0.5, 0.5],
11
+ "image_std": [0.5, 0.5, 0.5],
12
+ "max_side": 448,
13
+ "supported_modalities": ["text", "image"],
14
+ "unsupported_modalities": ["audio", "video", "image-generation", "video-generation"]
15
+ }
processing_dendro_omni.py ADDED
@@ -0,0 +1,267 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Image processor for Phillnet Mini Text-Vision.
2
+
3
+ The processor produces Qwen-compatible visual patch tensors for the retained
4
+ transplanted vision encoder and deliberately exposes no image or video synthesis
5
+ functionality.
6
+ """
7
+ from __future__ import annotations
8
+
9
+ import json
10
+ from pathlib import Path
11
+ from typing import Any, Sequence
12
+
13
+ import torch
14
+ from PIL import Image
15
+
16
+ from transformers import AutoTokenizer
17
+
18
+ try:
19
+ from transformers.feature_extraction_utils import BatchFeature
20
+ except Exception: # pragma: no cover
21
+ BatchFeature = dict # type: ignore[misc,assignment]
22
+
23
+
24
+ IMAGE_MARKER = "\ue000"
25
+
26
+
27
+ class DendroVisionProcessor:
28
+ """Prepare text and still-image inputs for Phillnet Mini Text-Vision.
29
+
30
+ This class supports text generation and image understanding only. It accepts
31
+ one conversation at a time, creates static two-frame visual patches required
32
+ by the retained vision tower, and expands one image placeholder per merged
33
+ visual token.
34
+ """
35
+
36
+ model_input_names = [
37
+ "input_ids",
38
+ "attention_mask",
39
+ "pixel_values",
40
+ "image_grid_thw",
41
+ "mm_token_type_ids",
42
+ ]
43
+
44
+ @classmethod
45
+ def register_for_auto_class(cls, auto_class: str = "AutoProcessor") -> None:
46
+ """Compatibility hook used by Transformers dynamic-module loading."""
47
+ cls._auto_class = str(auto_class)
48
+
49
+ def __init__(
50
+ self,
51
+ tokenizer: Any,
52
+ *,
53
+ image_token_id: int,
54
+ vision_start_token_id: int,
55
+ vision_end_token_id: int,
56
+ patch_size: int = 16,
57
+ temporal_patch_size: int = 2,
58
+ spatial_merge_size: int = 2,
59
+ image_mean: Sequence[float] = (0.5, 0.5, 0.5),
60
+ image_std: Sequence[float] = (0.5, 0.5, 0.5),
61
+ max_side: int = 448,
62
+ ) -> None:
63
+ self.tokenizer = tokenizer
64
+ self.image_token_id = int(image_token_id)
65
+ self.vision_start_token_id = int(vision_start_token_id)
66
+ self.vision_end_token_id = int(vision_end_token_id)
67
+ self.patch_size = int(patch_size)
68
+ self.temporal_patch_size = int(temporal_patch_size)
69
+ self.spatial_merge_size = int(spatial_merge_size)
70
+ self.image_mean = tuple(float(x) for x in image_mean)
71
+ self.image_std = tuple(float(x) for x in image_std)
72
+ self.max_side = int(max_side)
73
+ if self.patch_size < 1 or self.temporal_patch_size < 1 or self.spatial_merge_size < 1:
74
+ raise ValueError("Visual patch and merge dimensions must be positive")
75
+
76
+ @classmethod
77
+ def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs: Any) -> "DendroVisionProcessor":
78
+ root = Path(pretrained_model_name_or_path)
79
+ config = json.loads((root / "config.json").read_text(encoding="utf-8"))
80
+ preprocessing = json.loads((root / "preprocessor_config.json").read_text(encoding="utf-8"))
81
+ tokenizer_kwargs = dict(kwargs)
82
+ max_side = int(tokenizer_kwargs.pop("max_side", preprocessing.get("max_side", 448)))
83
+ tokenizer_kwargs.pop("trust_remote_code", None)
84
+ tokenizer_kwargs.pop("_from_auto", None)
85
+ tokenizer = AutoTokenizer.from_pretrained(str(root), trust_remote_code=False, **tokenizer_kwargs)
86
+ return cls(
87
+ tokenizer,
88
+ image_token_id=int(config["image_token_id"]),
89
+ vision_start_token_id=int(config["vision_start_token_id"]),
90
+ vision_end_token_id=int(config["vision_end_token_id"]),
91
+ patch_size=int(preprocessing.get("patch_size", 16)),
92
+ temporal_patch_size=int(preprocessing.get("temporal_patch_size", 2)),
93
+ spatial_merge_size=int(preprocessing.get("merge_size", 2)),
94
+ image_mean=preprocessing.get("image_mean", (0.5, 0.5, 0.5)),
95
+ image_std=preprocessing.get("image_std", (0.5, 0.5, 0.5)),
96
+ max_side=max_side,
97
+ )
98
+
99
+ def _encode_text(self, value: str) -> list[int]:
100
+ return list(self.tokenizer.encode(str(value), add_special_tokens=False))
101
+
102
+ @staticmethod
103
+ def _to_pil(image: Any) -> Image.Image:
104
+ if isinstance(image, Image.Image):
105
+ return image.convert("RGB")
106
+ if isinstance(image, (str, Path)):
107
+ with Image.open(image) as opened:
108
+ return opened.convert("RGB")
109
+ if torch.is_tensor(image):
110
+ value = image.detach().cpu().float()
111
+ if value.ndim == 4 and value.shape[0] == 1:
112
+ value = value[0]
113
+ if value.ndim != 3:
114
+ raise TypeError("Image tensor must be [C,H,W] or [1,C,H,W]")
115
+ if value.shape[0] in {1, 3, 4}:
116
+ value = value[:3].permute(1, 2, 0)
117
+ value = value.clamp(0, 1).mul(255).byte().numpy()
118
+ return Image.fromarray(value).convert("RGB")
119
+ try:
120
+ import numpy as np
121
+ value = np.asarray(image)
122
+ if value.ndim == 3:
123
+ if value.dtype != np.uint8:
124
+ value = (value.clip(0, 1) * 255 if float(value.max()) <= 1 else value.clip(0, 255)).astype(np.uint8)
125
+ return Image.fromarray(value).convert("RGB")
126
+ except Exception as error: # pragma: no cover
127
+ raise TypeError("Unsupported image input") from error
128
+ raise TypeError("Unsupported image input")
129
+
130
+ def _resize(self, image: Image.Image) -> Image.Image:
131
+ unit = self.patch_size * self.spatial_merge_size
132
+ width, height = image.size
133
+ if max(width, height) > self.max_side:
134
+ scale = self.max_side / max(width, height)
135
+ width, height = round(width * scale), round(height * scale)
136
+ width = max(unit, round(width / unit) * unit)
137
+ height = max(unit, round(height / unit) * unit)
138
+ return image.resize((width, height), Image.Resampling.BICUBIC)
139
+
140
+ def _patchify(self, image: Any) -> tuple[torch.Tensor, torch.Tensor, int]:
141
+ prepared = self._resize(self._to_pil(image))
142
+ try:
143
+ import numpy as np
144
+ values = torch.from_numpy(np.asarray(prepared).copy()).permute(2, 0, 1).float().div_(255.0)
145
+ except Exception as error: # pragma: no cover
146
+ raise RuntimeError("NumPy is required for image preprocessing") from error
147
+ mean = torch.tensor(self.image_mean).view(3, 1, 1)
148
+ std = torch.tensor(self.image_std).view(3, 1, 1)
149
+ values = (values - mean) / std
150
+ channels, height, width = values.shape
151
+ patch = self.patch_size
152
+ temporal = self.temporal_patch_size
153
+ grid_h, grid_w = height // patch, width // patch
154
+ frames = values.unsqueeze(0).repeat(temporal, 1, 1, 1)
155
+ blocks = frames.reshape(1, temporal, channels, grid_h, patch, grid_w, patch)
156
+ blocks = blocks.permute(0, 3, 5, 2, 1, 4, 6).reshape(-1, channels * temporal * patch * patch)
157
+ grid = torch.tensor([[1, grid_h, grid_w]], dtype=torch.long)
158
+ placeholder_count = grid_h // self.spatial_merge_size * (grid_w // self.spatial_merge_size)
159
+ return blocks, grid, int(placeholder_count)
160
+
161
+ def _build_single(self, text: str, images: Sequence[Any] | None) -> dict[str, torch.Tensor]:
162
+ image_list = list(images or [])
163
+ if IMAGE_MARKER not in text and image_list:
164
+ text = text + IMAGE_MARKER * len(image_list)
165
+ chunks = text.split(IMAGE_MARKER)
166
+ if len(chunks) != len(image_list) + 1:
167
+ raise ValueError("Image markers must match the number of supplied images")
168
+ ids: list[int] = []
169
+ types: list[int] = []
170
+ pixel_blocks: list[torch.Tensor] = []
171
+ grids: list[torch.Tensor] = []
172
+ for index, chunk in enumerate(chunks):
173
+ text_ids = self._encode_text(chunk)
174
+ ids.extend(text_ids)
175
+ types.extend([0] * len(text_ids))
176
+ if index == len(image_list):
177
+ continue
178
+ blocks, grid, count = self._patchify(image_list[index])
179
+ ids.append(self.vision_start_token_id)
180
+ types.append(0)
181
+ ids.extend([self.image_token_id] * count)
182
+ types.extend([1] * count)
183
+ ids.append(self.vision_end_token_id)
184
+ types.append(0)
185
+ pixel_blocks.append(blocks)
186
+ grids.append(grid)
187
+ output: dict[str, torch.Tensor] = {
188
+ "input_ids": torch.tensor([ids], dtype=torch.long),
189
+ "attention_mask": torch.ones((1, len(ids)), dtype=torch.long),
190
+ }
191
+ if pixel_blocks:
192
+ output["pixel_values"] = torch.cat(pixel_blocks, dim=0)
193
+ output["image_grid_thw"] = torch.cat(grids, dim=0)
194
+ output["mm_token_type_ids"] = torch.tensor([types], dtype=torch.long)
195
+ return output
196
+
197
+ def __call__(
198
+ self,
199
+ text: str | Sequence[str],
200
+ *,
201
+ images: Any | Sequence[Any] | None = None,
202
+ return_tensors: str | None = "pt",
203
+ **_: Any,
204
+ ) -> Any:
205
+ if isinstance(text, Sequence) and not isinstance(text, str):
206
+ if len(text) != 1:
207
+ raise ValueError("DendroVisionProcessor currently accepts a single conversation per call")
208
+ text = text[0]
209
+ if images is None:
210
+ image_list: list[Any] = []
211
+ elif isinstance(images, (str, Path, Image.Image)) or torch.is_tensor(images):
212
+ image_list = [images]
213
+ else:
214
+ image_list = list(images)
215
+ encoded = self._build_single(str(text), image_list)
216
+ if return_tensors not in {None, "pt"}:
217
+ raise ValueError("Only return_tensors='pt' is supported")
218
+ return BatchFeature(data=encoded, tensor_type="pt") if BatchFeature is not dict else encoded
219
+
220
+ def apply_chat_template(
221
+ self,
222
+ messages: Sequence[dict[str, Any]],
223
+ *,
224
+ tokenize: bool = True,
225
+ add_generation_prompt: bool = True,
226
+ enable_thinking: bool = False,
227
+ return_dict: bool = True,
228
+ return_tensors: str | None = "pt",
229
+ **_: Any,
230
+ ) -> Any:
231
+ parts: list[str] = []
232
+ images: list[Any] = []
233
+ for message in messages:
234
+ role = str(message.get("role", "user"))
235
+ parts.append(f"<|im_start|>{role}\n")
236
+ content = message.get("content", "")
237
+ if isinstance(content, str):
238
+ parts.append(content)
239
+ else:
240
+ for item in content:
241
+ item_type = item.get("type") if isinstance(item, dict) else None
242
+ if item_type == "text":
243
+ parts.append(str(item.get("text", "")))
244
+ elif item_type == "image":
245
+ image = item.get("image", item.get("image_url"))
246
+ if image is None:
247
+ raise ValueError("Image content must provide an image object")
248
+ images.append(image)
249
+ parts.append(IMAGE_MARKER)
250
+ else:
251
+ raise ValueError(f"Unsupported chat content type: {item_type!r}")
252
+ parts.append("<|im_end|>\n")
253
+ if add_generation_prompt:
254
+ parts.append("<|im_start|>assistant\n")
255
+ parts.append("<think>\n" if enable_thinking else "<think>\n\n</think>\n\n")
256
+ prompt = "".join(parts)
257
+ if not tokenize:
258
+ return prompt
259
+ encoded = self(prompt, images=images, return_tensors=return_tensors)
260
+ return encoded if return_dict else encoded["input_ids"]
261
+
262
+ def save_pretrained(self, save_directory: str | Path, **_: Any) -> tuple[str]:
263
+ root = Path(save_directory)
264
+ root.mkdir(parents=True, exist_ok=True)
265
+ path = root / "preprocessor_config.json"
266
+ path.write_text(json.dumps({"processor_class": "DendroVisionProcessor"}, indent=2) + "\n", encoding="utf-8")
267
+ return (str(path),)
requirements-server.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ fastapi>=0.115
2
+ uvicorn[standard]>=0.30
3
+ pydantic>=2.0
server.py ADDED
@@ -0,0 +1,226 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Production-oriented text-and-still-image API for Phillnet Mini Text-Vision.
2
+
3
+ The service intentionally exposes only chat completion and visual-question-answering
4
+ workflows. SDXL, image/video synthesis, audio, agents, tools, and remote image URL
5
+ fetching are outside this deployment surface.
6
+ """
7
+ from __future__ import annotations
8
+
9
+ import base64
10
+ import hmac
11
+ import io
12
+ import os
13
+ import threading
14
+ import time
15
+ import uuid
16
+ from contextlib import asynccontextmanager
17
+ from pathlib import Path
18
+ from typing import Any
19
+
20
+ import torch
21
+ from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
22
+ from fastapi.middleware.cors import CORSMiddleware
23
+ from fastapi.responses import JSONResponse
24
+ from PIL import Image, UnidentifiedImageError
25
+ from pydantic import BaseModel, Field
26
+ from transformers import AutoModelForCausalLM, AutoProcessor
27
+
28
+ MODEL_DIR = Path(os.getenv("MODEL_DIR", Path(__file__).resolve().parent))
29
+ API_KEY = os.getenv("PHILLNET_API_KEY", "")
30
+ MAX_IMAGE_BYTES = int(os.getenv("PHILLNET_MAX_IMAGE_BYTES", str(10 * 1024 * 1024)))
31
+ MAX_IMAGE_PIXELS = int(os.getenv("PHILLNET_MAX_IMAGE_PIXELS", str(24_000_000)))
32
+ MAX_REQUEST_BYTES = int(os.getenv("PHILLNET_MAX_REQUEST_BYTES", str(12 * 1024 * 1024)))
33
+ CORS_ORIGINS = [origin.strip() for origin in os.getenv("PHILLNET_CORS_ORIGINS", "").split(",") if origin.strip()]
34
+
35
+ Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
36
+ MODEL: Any | None = None
37
+ PROCESSOR: Any | None = None
38
+ GENERATION_LOCK = threading.Lock()
39
+
40
+
41
+ class ImageContent(BaseModel):
42
+ type: str
43
+ text: str | None = Field(default=None, max_length=32_000)
44
+ image_base64: str | None = None
45
+
46
+
47
+ class Message(BaseModel):
48
+ role: str = Field(pattern="^(system|user|assistant)$")
49
+ content: str | list[ImageContent]
50
+
51
+
52
+ class ChatRequest(BaseModel):
53
+ model: str = "phillnet-mini-text-vision"
54
+ messages: list[Message] = Field(min_length=1, max_length=32)
55
+ max_tokens: int = Field(default=8192, ge=1, le=8192)
56
+ temperature: float = Field(default=0.0, ge=0.0, le=2.0)
57
+ reasoning_effort: str = Field(default="max", pattern="^(direct|low|medium|high|max)$")
58
+
59
+
60
+ def require_api_key(
61
+ authorization: str | None = Header(default=None),
62
+ x_api_key: str | None = Header(default=None),
63
+ ) -> None:
64
+ """Enforce an API key when PHILLNET_API_KEY is configured.
65
+
66
+ Local development remains frictionless when the environment variable is empty.
67
+ Production compose configuration supplies a non-empty secret by default.
68
+ """
69
+ if not API_KEY:
70
+ return
71
+ candidate = x_api_key or ""
72
+ if authorization and authorization.lower().startswith("bearer "):
73
+ candidate = authorization[7:].strip()
74
+ if not candidate or not hmac.compare_digest(candidate, API_KEY):
75
+ raise HTTPException(
76
+ status_code=status.HTTP_401_UNAUTHORIZED,
77
+ detail="Valid API credentials are required.",
78
+ headers={"WWW-Authenticate": "Bearer"},
79
+ )
80
+
81
+
82
+ def decode_image(encoded: str) -> Image.Image:
83
+ try:
84
+ raw = encoded.split(",", 1)[1] if encoded.startswith("data:") else encoded
85
+ # Base64 expands bytes by roughly 4/3. Guard before decoding a large body.
86
+ if len(raw) > ((MAX_IMAGE_BYTES * 4) // 3) + 8:
87
+ raise HTTPException(413, f"image_base64 exceeds the {MAX_IMAGE_BYTES}-byte limit")
88
+ payload = base64.b64decode(raw, validate=True)
89
+ if len(payload) > MAX_IMAGE_BYTES:
90
+ raise HTTPException(413, f"image_base64 exceeds the {MAX_IMAGE_BYTES}-byte limit")
91
+ image = Image.open(io.BytesIO(payload))
92
+ image.load()
93
+ return image.convert("RGB")
94
+ except HTTPException:
95
+ raise
96
+ except (ValueError, UnidentifiedImageError, OSError, Image.DecompressionBombError) as error:
97
+ raise HTTPException(400, "image_base64 must be a valid, safe base64-encoded image") from error
98
+
99
+
100
+ def make_contents(messages: list[Message]) -> list[dict[str, Any]]:
101
+ result: list[dict[str, Any]] = []
102
+ image_count = 0
103
+ for message in messages:
104
+ if isinstance(message.content, str):
105
+ if len(message.content) > 32_000:
106
+ raise HTTPException(400, "A text message may not exceed 32,000 characters")
107
+ content: str | list[dict[str, Any]] = message.content
108
+ else:
109
+ content = []
110
+ for item in message.content:
111
+ if item.type == "text":
112
+ content.append({"type": "text", "text": item.text or ""})
113
+ elif item.type == "image":
114
+ image_count += 1
115
+ if image_count > 4:
116
+ raise HTTPException(400, "A request may contain at most four images")
117
+ if not item.image_base64:
118
+ raise HTTPException(400, "image content requires image_base64")
119
+ content.append({"type": "image", "image": decode_image(item.image_base64)})
120
+ else:
121
+ raise HTTPException(400, f"Unsupported content type: {item.type!r}. Only text and image are supported.")
122
+ result.append({"role": message.role, "content": content})
123
+ return result
124
+
125
+
126
+ @asynccontextmanager
127
+ async def lifespan(_app: FastAPI):
128
+ global MODEL, PROCESSOR
129
+ PROCESSOR = AutoProcessor.from_pretrained(str(MODEL_DIR), trust_remote_code=True)
130
+ MODEL = AutoModelForCausalLM.from_pretrained(
131
+ str(MODEL_DIR),
132
+ trust_remote_code=True,
133
+ dtype=torch.bfloat16,
134
+ low_cpu_mem_usage=True,
135
+ ).eval()
136
+ yield
137
+ MODEL = None
138
+ PROCESSOR = None
139
+
140
+
141
+ app = FastAPI(
142
+ title="Phillnet Mini Text-Vision",
143
+ version="1.1.0",
144
+ description="Text generation and still-image understanding only. SDXL, image generation, video generation, audio, tools, and agent runtimes are disabled.",
145
+ lifespan=lifespan,
146
+ )
147
+
148
+ if CORS_ORIGINS:
149
+ app.add_middleware(
150
+ CORSMiddleware,
151
+ allow_origins=CORS_ORIGINS,
152
+ allow_credentials=False,
153
+ allow_methods=["GET", "POST"],
154
+ allow_headers=["Authorization", "Content-Type", "X-API-Key"],
155
+ max_age=600,
156
+ )
157
+
158
+
159
+ @app.middleware("http")
160
+ async def enforce_request_limit(request: Request, call_next: Any) -> Any:
161
+ content_length = request.headers.get("content-length")
162
+ if content_length and int(content_length) > MAX_REQUEST_BYTES:
163
+ return JSONResponse(status_code=413, content={"detail": "Request body exceeds configured size limit"})
164
+ return await call_next(request)
165
+
166
+
167
+ @app.get("/health")
168
+ def health() -> dict[str, Any]:
169
+ ready = MODEL is not None and PROCESSOR is not None
170
+ return {
171
+ "status": "ok" if ready else "loading",
172
+ "ready": ready,
173
+ "service": "phillnet-mini-text-vision",
174
+ "version": app.version,
175
+ "capabilities": ["text-generation", "image-understanding"],
176
+ "disabled": ["image-generation", "video-generation", "audio", "tools", "agents"],
177
+ }
178
+
179
+
180
+ @app.get("/ready")
181
+ def ready() -> dict[str, bool]:
182
+ if MODEL is None or PROCESSOR is None:
183
+ raise HTTPException(503, "Model is still loading")
184
+ return {"ready": True}
185
+
186
+
187
+ @app.post("/v1/chat/completions", dependencies=[Depends(require_api_key)])
188
+ def chat_completions(request: ChatRequest) -> dict[str, Any]:
189
+ if MODEL is None or PROCESSOR is None:
190
+ raise HTTPException(503, "Model is still loading")
191
+ content = make_contents(request.messages)
192
+ encoded = PROCESSOR.apply_chat_template(
193
+ content,
194
+ tokenize=True,
195
+ add_generation_prompt=True,
196
+ return_dict=True,
197
+ return_tensors="pt",
198
+ )
199
+ device = next(MODEL.parameters()).device
200
+ encoded = {key: value.to(device) if torch.is_tensor(value) else value for key, value in dict(encoded).items()}
201
+ prompt_length = int(encoded["input_ids"].shape[1])
202
+ generation_kwargs: dict[str, Any] = {
203
+ "max_new_tokens": request.max_tokens,
204
+ "do_sample": request.temperature > 0.0,
205
+ "use_cache": True,
206
+ "reasoning_effort": request.reasoning_effort,
207
+ }
208
+ if request.temperature > 0.0:
209
+ generation_kwargs["temperature"] = request.temperature
210
+ started = time.perf_counter()
211
+ # A single local model instance should perform one generation at a time to
212
+ # prevent concurrent high-context calls from overcommitting model memory.
213
+ with GENERATION_LOCK, torch.inference_mode():
214
+ output = MODEL.generate(**encoded, **generation_kwargs)
215
+ completion_ids = output[0, prompt_length:].detach().cpu()
216
+ text = PROCESSOR.tokenizer.decode(completion_ids, skip_special_tokens=True)
217
+ completion_tokens = int(completion_ids.numel())
218
+ return {
219
+ "id": f"chatcmpl-{uuid.uuid4().hex}",
220
+ "object": "chat.completion",
221
+ "created": int(time.time()),
222
+ "model": "phillnet-mini-text-vision",
223
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
224
+ "usage": {"prompt_tokens": prompt_length, "completion_tokens": completion_tokens, "total_tokens": prompt_length + completion_tokens},
225
+ "elapsed_seconds": round(time.perf_counter() - started, 3),
226
+ }
tokenization_dendro_omni.py ADDED
@@ -0,0 +1,334 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Deterministic UTF-8 byte tokenizer for Dendro Omni."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ from pathlib import Path
7
+ from typing import Any, Iterable, Sequence
8
+
9
+ import torch
10
+
11
+ try: # pragma: no cover - optional Hugging Face branch.
12
+ from transformers import PreTrainedTokenizer
13
+
14
+ _HF_TOKENIZER = True
15
+ except Exception: # pragma: no cover - fallback is covered.
16
+ PreTrainedTokenizer = object # type: ignore[assignment,misc]
17
+ _HF_TOKENIZER = False
18
+
19
+
20
+ PAD_TOKEN = "<pad>"
21
+ BOS_TOKEN = "<bos>"
22
+ EOS_TOKEN = "<eos>"
23
+ BYTE_OFFSET = 3
24
+
25
+
26
+ def _byte_token(value: int) -> str:
27
+ return f"<0x{value:02X}>"
28
+
29
+
30
+ def _token_byte(token: str) -> int | None:
31
+ if len(token) == 6 and token.startswith("<0x") and token.endswith(">"):
32
+ try:
33
+ value = int(token[3:5], 16)
34
+ except ValueError:
35
+ return None
36
+ return value if 0 <= value <= 255 else None
37
+ return None
38
+
39
+
40
+ class _ByteTokenizerLogic:
41
+ pad_token_id = 0
42
+ bos_token_id = 1
43
+ eos_token_id = 2
44
+ byte_offset = BYTE_OFFSET
45
+ model_input_names = ["input_ids", "attention_mask"]
46
+
47
+ @property
48
+ def vocab_size(self) -> int:
49
+ return self.byte_offset + 256
50
+
51
+ def raw_encode(self, text: str, *, add_bos: bool = False, add_eos: bool = False) -> list[int]:
52
+ ids: list[int] = []
53
+ if add_bos:
54
+ ids.append(self.bos_token_id)
55
+ ids.extend(byte + self.byte_offset for byte in text.encode("utf-8"))
56
+ if add_eos:
57
+ ids.append(self.eos_token_id)
58
+ return ids
59
+
60
+ def raw_decode(self, token_ids: Iterable[int], *, skip_special_tokens: bool = True) -> str:
61
+ data = bytearray()
62
+ literal: list[str] = []
63
+ for item in token_ids:
64
+ token_id = int(item)
65
+ if token_id in (self.pad_token_id, self.bos_token_id, self.eos_token_id):
66
+ if skip_special_tokens:
67
+ continue
68
+ literal.append({0: PAD_TOKEN, 1: BOS_TOKEN, 2: EOS_TOKEN}[token_id])
69
+ continue
70
+ value = token_id - self.byte_offset
71
+ if 0 <= value <= 255:
72
+ if literal:
73
+ data.extend("".join(literal).encode("utf-8"))
74
+ literal.clear()
75
+ data.append(value)
76
+ if literal:
77
+ data.extend("".join(literal).encode("utf-8"))
78
+ return data.decode("utf-8", errors="replace")
79
+
80
+
81
+ if _HF_TOKENIZER:
82
+ class DendroByteTokenizer(_ByteTokenizerLogic, PreTrainedTokenizer):
83
+ """Hugging Face tokenizer with lossless byte coverage and no learned files."""
84
+
85
+ vocab_files_names = {"vocab_file": "dendro_vocab.json"}
86
+ model_input_names = ["input_ids", "attention_mask"]
87
+
88
+ def __init__(
89
+ self,
90
+ vocab_file: str | None = None,
91
+ *,
92
+ model_max_length: int = 131_072,
93
+ **kwargs: Any,
94
+ ) -> None:
95
+ del vocab_file
96
+ # ``from_pretrained`` restores these values through ``kwargs``. Use
97
+ # defaults instead of passing a second explicit copy so this remains
98
+ # compatible with both legacy tokenizers and the Transformers 5
99
+ # Python backend.
100
+ kwargs.setdefault("pad_token", PAD_TOKEN)
101
+ kwargs.setdefault("bos_token", BOS_TOKEN)
102
+ kwargs.setdefault("eos_token", EOS_TOKEN)
103
+ kwargs.setdefault("model_max_length", model_max_length)
104
+ kwargs.setdefault("clean_up_tokenization_spaces", False)
105
+ super().__init__(**kwargs)
106
+
107
+ def _tokenize(self, text: str, **kwargs: Any) -> list[str]:
108
+ del kwargs
109
+ return [_byte_token(byte) for byte in text.encode("utf-8")]
110
+
111
+ def _convert_token_to_id(self, token: str) -> int:
112
+ if token == PAD_TOKEN:
113
+ return self.pad_token_id
114
+ if token == BOS_TOKEN:
115
+ return self.bos_token_id
116
+ if token == EOS_TOKEN:
117
+ return self.eos_token_id
118
+ value = _token_byte(token)
119
+ return self.eos_token_id if value is None else value + self.byte_offset
120
+
121
+ def _convert_id_to_token(self, index: int) -> str:
122
+ index = int(index)
123
+ if index == self.pad_token_id:
124
+ return PAD_TOKEN
125
+ if index == self.bos_token_id:
126
+ return BOS_TOKEN
127
+ if index == self.eos_token_id:
128
+ return EOS_TOKEN
129
+ value = index - self.byte_offset
130
+ return _byte_token(value) if 0 <= value <= 255 else EOS_TOKEN
131
+
132
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
133
+ data = bytearray()
134
+ text_parts: list[str] = []
135
+ for token in tokens:
136
+ value = _token_byte(token)
137
+ if value is None:
138
+ if data:
139
+ text_parts.append(data.decode("utf-8", errors="replace"))
140
+ data.clear()
141
+ if token not in {PAD_TOKEN, BOS_TOKEN, EOS_TOKEN}:
142
+ text_parts.append(token)
143
+ else:
144
+ data.append(value)
145
+ if data:
146
+ text_parts.append(data.decode("utf-8", errors="replace"))
147
+ return "".join(text_parts)
148
+
149
+ def get_vocab(self) -> dict[str, int]:
150
+ vocab = {PAD_TOKEN: 0, BOS_TOKEN: 1, EOS_TOKEN: 2}
151
+ vocab.update({_byte_token(value): value + self.byte_offset for value in range(256)})
152
+ return vocab
153
+
154
+ def build_inputs_with_special_tokens(
155
+ self,
156
+ token_ids_0: list[int],
157
+ token_ids_1: list[int] | None = None,
158
+ ) -> list[int]:
159
+ result = [self.bos_token_id, *token_ids_0, self.eos_token_id]
160
+ if token_ids_1 is not None:
161
+ result.extend([*token_ids_1, self.eos_token_id])
162
+ return result
163
+
164
+ def get_special_tokens_mask(
165
+ self,
166
+ token_ids_0: list[int],
167
+ token_ids_1: list[int] | None = None,
168
+ already_has_special_tokens: bool = False,
169
+ ) -> list[int]:
170
+ if already_has_special_tokens:
171
+ return [int(token in {0, 1, 2}) for token in token_ids_0]
172
+ mask = [1] + [0] * len(token_ids_0) + [1]
173
+ if token_ids_1 is not None:
174
+ mask += [0] * len(token_ids_1) + [1]
175
+ return mask
176
+
177
+ def save_vocabulary(
178
+ self,
179
+ save_directory: str,
180
+ filename_prefix: str | None = None,
181
+ ) -> tuple[str]:
182
+ directory = Path(save_directory)
183
+ directory.mkdir(parents=True, exist_ok=True)
184
+ name = f"{filename_prefix + '-' if filename_prefix else ''}dendro_vocab.json"
185
+ path = directory / name
186
+ path.write_text(json.dumps(self.get_vocab(), indent=2, sort_keys=True), encoding="utf-8")
187
+ return (str(path),)
188
+
189
+ def encode_bytes(self, text: str, *, add_bos: bool = False, add_eos: bool = False) -> list[int]:
190
+ return self.raw_encode(text, add_bos=add_bos, add_eos=add_eos)
191
+
192
+ def decode_bytes(self, ids: Iterable[int], *, skip_special_tokens: bool = True) -> str:
193
+ return self.raw_decode(ids, skip_special_tokens=skip_special_tokens)
194
+
195
+ else:
196
+ class DendroByteTokenizer(_ByteTokenizerLogic):
197
+ """PyTorch-only fallback implementing the common tokenizer call surface."""
198
+
199
+ def __init__(
200
+ self,
201
+ vocab_file: str | None = None,
202
+ *,
203
+ model_max_length: int = 131_072,
204
+ padding_side: str = "right",
205
+ **_: Any,
206
+ ) -> None:
207
+ del vocab_file
208
+ self.model_max_length = int(model_max_length)
209
+ self.padding_side = str(padding_side)
210
+ self.pad_token = PAD_TOKEN
211
+ self.bos_token = BOS_TOKEN
212
+ self.eos_token = EOS_TOKEN
213
+
214
+ def get_vocab(self) -> dict[str, int]:
215
+ vocab = {PAD_TOKEN: 0, BOS_TOKEN: 1, EOS_TOKEN: 2}
216
+ vocab.update({_byte_token(value): value + self.byte_offset for value in range(256)})
217
+ return vocab
218
+
219
+ def encode(
220
+ self,
221
+ text: str,
222
+ *,
223
+ add_special_tokens: bool = True,
224
+ add_bos: bool | None = None,
225
+ add_eos: bool | None = None,
226
+ **_: Any,
227
+ ) -> list[int]:
228
+ return self.raw_encode(
229
+ text,
230
+ add_bos=add_special_tokens if add_bos is None else add_bos,
231
+ add_eos=add_special_tokens if add_eos is None else add_eos,
232
+ )
233
+
234
+ def decode(
235
+ self,
236
+ token_ids: Iterable[int] | torch.Tensor,
237
+ *,
238
+ skip_special_tokens: bool = True,
239
+ **_: Any,
240
+ ) -> str:
241
+ if torch.is_tensor(token_ids):
242
+ token_ids = token_ids.detach().cpu().tolist()
243
+ return self.raw_decode(token_ids, skip_special_tokens=skip_special_tokens)
244
+
245
+ def batch_decode(
246
+ self,
247
+ sequences: Sequence[Sequence[int] | torch.Tensor],
248
+ *,
249
+ skip_special_tokens: bool = True,
250
+ **kwargs: Any,
251
+ ) -> list[str]:
252
+ return [self.decode(row, skip_special_tokens=skip_special_tokens, **kwargs) for row in sequences]
253
+
254
+ def __call__(
255
+ self,
256
+ text: str | Sequence[str],
257
+ *,
258
+ add_special_tokens: bool = True,
259
+ padding: bool | str = False,
260
+ truncation: bool = False,
261
+ max_length: int | None = None,
262
+ return_tensors: str | None = None,
263
+ **_: Any,
264
+ ) -> dict[str, Any]:
265
+ texts = [text] if isinstance(text, str) else list(text)
266
+ rows = [self.encode(item, add_special_tokens=add_special_tokens) for item in texts]
267
+ limit = self.model_max_length if max_length is None else int(max_length)
268
+ if truncation:
269
+ rows = [row[:limit] for row in rows]
270
+ target = max((len(row) for row in rows), default=0) if padding else None
271
+ if padding == "max_length":
272
+ target = limit
273
+ masks: list[list[int]] = []
274
+ if target is not None:
275
+ padded: list[list[int]] = []
276
+ for row in rows:
277
+ amount = max(0, target - len(row))
278
+ if self.padding_side == "left":
279
+ padded.append([self.pad_token_id] * amount + row[:target])
280
+ masks.append([0] * amount + [1] * min(len(row), target))
281
+ else:
282
+ padded.append(row[:target] + [self.pad_token_id] * amount)
283
+ masks.append([1] * min(len(row), target) + [0] * amount)
284
+ rows = padded
285
+ else:
286
+ masks = [[1] * len(row) for row in rows]
287
+ result: dict[str, Any] = {"input_ids": rows, "attention_mask": masks}
288
+ if return_tensors is not None:
289
+ if return_tensors != "pt":
290
+ raise ValueError("The local tokenizer supports return_tensors='pt' only")
291
+ if not padding and len({len(row) for row in rows}) > 1:
292
+ raise ValueError("Batch tensor output requires padding")
293
+ result = {key: torch.tensor(value, dtype=torch.long) for key, value in result.items()}
294
+ if isinstance(text, str) and return_tensors is None:
295
+ result = {key: value[0] for key, value in result.items()}
296
+ return result
297
+
298
+ def save_pretrained(self, save_directory: str | Path, **_: Any) -> tuple[str, ...]:
299
+ directory = Path(save_directory)
300
+ directory.mkdir(parents=True, exist_ok=True)
301
+ vocab_path = directory / "dendro_vocab.json"
302
+ config_path = directory / "tokenizer_config.json"
303
+ special_path = directory / "special_tokens_map.json"
304
+ vocab_path.write_text(json.dumps(self.get_vocab(), indent=2, sort_keys=True), encoding="utf-8")
305
+ config_path.write_text(
306
+ json.dumps(
307
+ {
308
+ "tokenizer_class": "DendroByteTokenizer",
309
+ "model_max_length": self.model_max_length,
310
+ "padding_side": self.padding_side,
311
+ "auto_map": {"AutoTokenizer": ["tokenization_dendro_omni.DendroByteTokenizer", None]},
312
+ },
313
+ indent=2,
314
+ ),
315
+ encoding="utf-8",
316
+ )
317
+ special_path.write_text(
318
+ json.dumps(
319
+ {"pad_token": PAD_TOKEN, "bos_token": BOS_TOKEN, "eos_token": EOS_TOKEN},
320
+ indent=2,
321
+ ),
322
+ encoding="utf-8",
323
+ )
324
+ return str(vocab_path), str(config_path), str(special_path)
325
+
326
+ @classmethod
327
+ def from_pretrained(cls, path: str | Path, **kwargs: Any) -> "DendroByteTokenizer":
328
+ directory = Path(path)
329
+ config_path = directory / "tokenizer_config.json"
330
+ if config_path.exists():
331
+ config = json.loads(config_path.read_text(encoding="utf-8"))
332
+ kwargs.setdefault("model_max_length", config.get("model_max_length", 131_072))
333
+ kwargs.setdefault("padding_side", config.get("padding_side", "right"))
334
+ return cls(vocab_file=str(directory / "dendro_vocab.json"), **kwargs)
tokenizer_config.json ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "audio_bos_token": "<|audio_start|>",
4
+ "audio_eos_token": "<|audio_end|>",
5
+ "audio_token": "<|audio_pad|>",
6
+ "backend": "tokenizers",
7
+ "bos_token": null,
8
+ "clean_up_tokenization_spaces": false,
9
+ "eos_token": "<|im_end|>",
10
+ "errors": "replace",
11
+ "image_token": "<|image_pad|>",
12
+ "is_local": true,
13
+ "local_files_only": true,
14
+ "model_max_length": 262144,
15
+ "model_specific_special_tokens": {
16
+ "audio_bos_token": "<|audio_start|>",
17
+ "audio_eos_token": "<|audio_end|>",
18
+ "audio_token": "<|audio_pad|>",
19
+ "image_token": "<|image_pad|>",
20
+ "video_token": "<|video_pad|>",
21
+ "vision_bos_token": "<|vision_start|>",
22
+ "vision_eos_token": "<|vision_end|>"
23
+ },
24
+ "pad_token": "<|endoftext|>",
25
+ "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
26
+ "split_special_tokens": false,
27
+ "tokenizer_class": "Qwen2Tokenizer",
28
+ "unk_token": null,
29
+ "video_token": "<|video_pad|>",
30
+ "vision_bos_token": "<|vision_start|>",
31
+ "vision_eos_token": "<|vision_end|>"
32
+ }
transplant.py ADDED
@@ -0,0 +1,1410 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Dendro-native exact-equivalence path for transplanted Qwen3.5 tensors.
2
+
3
+ This module contains no Qwen modules and registers no parameters. It evaluates
4
+ source-model operators directly from named slices of the one Dendro source
5
+ parameter. Dendro remains the owning architecture; the donor checkpoint is only
6
+ an initialization and exact-equivalence coordinate system.
7
+
8
+ The operator equations follow the Apache-2.0 Hugging Face Qwen3.5 reference
9
+ implementation, rewritten as parameterless Dendro functions so that transplant
10
+ weights never become a second model object.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import itertools
16
+ import math
17
+ from dataclasses import dataclass, field
18
+ from typing import Any, Mapping
19
+
20
+ import torch
21
+ from torch.nn import functional as F
22
+ from torch.utils.checkpoint import checkpoint
23
+
24
+ from ._source_bound import SourceBoundModule
25
+ from .accelerators import (
26
+ fla_causal_conv1d,
27
+ fla_causal_conv1d_update,
28
+ fla_chunk_gated_delta_rule,
29
+ fla_recurrent_gated_delta_rule,
30
+ )
31
+ from .source import DendroSourceLayer
32
+
33
+
34
+ def _can_use_fused_sdpa(
35
+ query: torch.Tensor,
36
+ key: torch.Tensor,
37
+ value: torch.Tensor,
38
+ *,
39
+ is_causal: bool,
40
+ enable_gqa: bool,
41
+ ) -> bool:
42
+ """Return whether PyTorch can run attention without a quadratic score tensor.
43
+
44
+ ``scaled_dot_product_attention`` silently falls back to its math backend when
45
+ a fused CUDA kernel is unavailable. That fallback still materializes the
46
+ complete attention matrix, so use PyTorch's backend probes before selecting
47
+ SDPA. Every probe is optional to retain compatibility with older PyTorch
48
+ releases and CPU-only installations.
49
+ """
50
+
51
+ if query.device.type != "cuda" or not hasattr(F, "scaled_dot_product_attention"):
52
+ return False
53
+ cuda_backends = getattr(torch.backends, "cuda", None)
54
+ params_type = getattr(cuda_backends, "SDPAParams", None)
55
+ if params_type is None:
56
+ return False
57
+ try:
58
+ params = params_type(query, key, value, None, 0.0, is_causal, enable_gqa)
59
+ except (RuntimeError, TypeError):
60
+ return False
61
+ for name in (
62
+ "can_use_flash_attention",
63
+ "can_use_efficient_attention",
64
+ "can_use_cudnn_attention",
65
+ ):
66
+ probe = getattr(cuda_backends, name, None)
67
+ if probe is None:
68
+ continue
69
+ try:
70
+ if bool(probe(params, False)):
71
+ return True
72
+ except (RuntimeError, TypeError):
73
+ continue
74
+ return False
75
+
76
+
77
+ def _run_fused_sdpa(
78
+ query: torch.Tensor,
79
+ key: torch.Tensor,
80
+ value: torch.Tensor,
81
+ *,
82
+ is_causal: bool,
83
+ enable_gqa: bool,
84
+ ) -> torch.Tensor | None:
85
+ """Run fused SDPA when supported, otherwise preserve the reference route."""
86
+
87
+ if not _can_use_fused_sdpa(
88
+ query,
89
+ key,
90
+ value,
91
+ is_causal=is_causal,
92
+ enable_gqa=enable_gqa,
93
+ ):
94
+ return None
95
+ try:
96
+ return F.scaled_dot_product_attention(
97
+ query,
98
+ key,
99
+ value,
100
+ dropout_p=0.0,
101
+ is_causal=is_causal,
102
+ enable_gqa=enable_gqa,
103
+ )
104
+ except (NotImplementedError, TypeError):
105
+ return None
106
+ except RuntimeError as error:
107
+ # Do not hide a real allocation failure behind a second, larger
108
+ # reference allocation. Other backend launch failures safely fall back.
109
+ if "out of memory" in str(error).lower():
110
+ raise
111
+ return None
112
+
113
+
114
+ def _torch_grouped_query_attention(
115
+ query: torch.Tensor,
116
+ key: torch.Tensor,
117
+ value: torch.Tensor,
118
+ *,
119
+ past_length: int,
120
+ attention_mask: torch.Tensor | None,
121
+ ) -> torch.Tensor:
122
+ """Reference GQA without physically repeating key and value heads.
123
+
124
+ Combining the repetition factor with the query-length dimension turns GQA
125
+ into one batched matrix multiplication per native KV head. This retains the
126
+ donor equation while avoiding a persistent or transient ``repeat_kv`` copy.
127
+ """
128
+
129
+ batch, num_heads, length, head_dim = query.shape
130
+ kv_heads = key.shape[1]
131
+ repetitions = num_heads // kv_heads
132
+ key_length = key.shape[-2]
133
+ grouped_query = query.reshape(batch * kv_heads, repetitions * length, head_dim)
134
+ grouped_key = key.reshape(batch * kv_heads, key_length, head_dim)
135
+ grouped_value = value.reshape(batch * kv_heads, key_length, head_dim)
136
+ scores = torch.bmm(grouped_query, grouped_key.transpose(1, 2))
137
+ scores = scores.reshape(batch, kv_heads, repetitions, length, key_length)
138
+ scores = scores * (head_dim**-0.5)
139
+
140
+ query_positions = past_length + torch.arange(length, device=scores.device)
141
+ key_positions = torch.arange(key_length, device=scores.device)
142
+ causal = key_positions.unsqueeze(0) > query_positions.unsqueeze(1)
143
+ scores = scores.masked_fill(causal[None, None, None], torch.finfo(scores.dtype).min)
144
+ if attention_mask is not None:
145
+ key_mask = ~attention_mask.to(device=scores.device, dtype=torch.bool)
146
+ scores = scores.masked_fill(
147
+ key_mask[:, None, None, None, :],
148
+ torch.finfo(scores.dtype).min,
149
+ )
150
+ probabilities = F.softmax(scores.float(), dim=-1).to(query.dtype)
151
+ output = torch.bmm(
152
+ probabilities.reshape(batch * kv_heads, repetitions * length, key_length),
153
+ grouped_value,
154
+ )
155
+ return output.reshape(batch, num_heads, length, head_dim)
156
+
157
+
158
+ @dataclass(slots=True)
159
+ class DendroTransplantOutput:
160
+ last_hidden_state: torch.Tensor
161
+ input_embeddings: torch.Tensor
162
+ position_ids: torch.Tensor
163
+ rope_deltas: torch.Tensor | None = None
164
+ cache: "DendroTransplantCache | None" = None
165
+
166
+
167
+ @dataclass(slots=True)
168
+ class DendroTransplantCache:
169
+ """Parameterless incremental state for the exact donor operators."""
170
+
171
+ position: int = 0
172
+ attention_mask: torch.Tensor | None = None
173
+ full_attention: dict[int, tuple[torch.Tensor, torch.Tensor]] = field(default_factory=dict)
174
+ linear_convolution: dict[int, torch.Tensor] = field(default_factory=dict)
175
+ linear_recurrent: dict[int, torch.Tensor] = field(default_factory=dict)
176
+
177
+ @property
178
+ def total_tokens_seen(self) -> int:
179
+ return int(self.position)
180
+
181
+ def get_seq_length(self, layer_idx: int = 0) -> int:
182
+ del layer_idx
183
+ return int(self.position)
184
+
185
+ def reorder(self, indices: torch.Tensor) -> "DendroTransplantCache":
186
+ if self.attention_mask is not None:
187
+ self.attention_mask = self.attention_mask.index_select(0, indices.to(self.attention_mask.device))
188
+ self.full_attention = {
189
+ layer: (
190
+ key.index_select(0, indices.to(key.device)),
191
+ value.index_select(0, indices.to(value.device)),
192
+ )
193
+ for layer, (key, value) in self.full_attention.items()
194
+ }
195
+ self.linear_convolution = {
196
+ layer: value.index_select(0, indices.to(value.device))
197
+ for layer, value in self.linear_convolution.items()
198
+ }
199
+ self.linear_recurrent = {
200
+ layer: value.index_select(0, indices.to(value.device))
201
+ for layer, value in self.linear_recurrent.items()
202
+ }
203
+ return self
204
+
205
+ def reorder_cache(self, indices: torch.Tensor) -> "DendroTransplantCache":
206
+ return self.reorder(indices)
207
+
208
+
209
+ def required_transplant_tensor_names(source_config: Mapping[str, Any]) -> set[str]:
210
+ """Return every donor tensor required by Dendro's exact text/vision path.
211
+
212
+ This is deliberately independent of a live model or source tensor so the root
213
+ transplant script can reject an incomplete checkpoint *before* allocating and
214
+ writing the packed one-source target. MTP tensors are preserved in the donor
215
+ bank when present, but are not required for ordinary causal-LM equivalence.
216
+ """
217
+
218
+ text_config: Mapping[str, Any] = source_config.get("text_config", source_config)
219
+ vision_config: Mapping[str, Any] = source_config.get("vision_config", {})
220
+ language_prefix = "model.language_model"
221
+ required: set[str] = {
222
+ f"{language_prefix}.embed_tokens.weight",
223
+ f"{language_prefix}.norm.weight",
224
+ }
225
+
226
+ layer_types = list(text_config.get("layer_types", []))
227
+ num_layers = int(text_config.get("num_hidden_layers", len(layer_types)))
228
+ full_attention_interval = int(text_config.get("full_attention_interval", 4))
229
+ attention_bias = bool(text_config.get("attention_bias", False))
230
+ for index in range(num_layers):
231
+ prefix = f"{language_prefix}.layers.{index}"
232
+ required.update(
233
+ {
234
+ f"{prefix}.input_layernorm.weight",
235
+ f"{prefix}.post_attention_layernorm.weight",
236
+ f"{prefix}.mlp.gate_proj.weight",
237
+ f"{prefix}.mlp.up_proj.weight",
238
+ f"{prefix}.mlp.down_proj.weight",
239
+ }
240
+ )
241
+ layer_type = layer_types[index] if layer_types else (
242
+ "full_attention" if (index + 1) % full_attention_interval == 0 else "linear_attention"
243
+ )
244
+ if layer_type == "full_attention":
245
+ attention = f"{prefix}.self_attn"
246
+ required.update(
247
+ {
248
+ f"{attention}.q_proj.weight",
249
+ f"{attention}.k_proj.weight",
250
+ f"{attention}.v_proj.weight",
251
+ f"{attention}.o_proj.weight",
252
+ f"{attention}.q_norm.weight",
253
+ f"{attention}.k_norm.weight",
254
+ }
255
+ )
256
+ if attention_bias:
257
+ required.update(
258
+ {
259
+ f"{attention}.q_proj.bias",
260
+ f"{attention}.k_proj.bias",
261
+ f"{attention}.v_proj.bias",
262
+ f"{attention}.o_proj.bias",
263
+ }
264
+ )
265
+ elif layer_type == "linear_attention":
266
+ attention = f"{prefix}.linear_attn"
267
+ required.update(
268
+ {
269
+ f"{attention}.in_proj_qkv.weight",
270
+ f"{attention}.in_proj_z.weight",
271
+ f"{attention}.in_proj_b.weight",
272
+ f"{attention}.in_proj_a.weight",
273
+ f"{attention}.conv1d.weight",
274
+ f"{attention}.dt_bias",
275
+ f"{attention}.A_log",
276
+ f"{attention}.norm.weight",
277
+ f"{attention}.out_proj.weight",
278
+ }
279
+ )
280
+ else:
281
+ raise ValueError(f"unsupported Qwen3.5 donor layer type at index {index}: {layer_type!r}")
282
+
283
+ if vision_config:
284
+ required.update(
285
+ {
286
+ "model.visual.patch_embed.proj.weight",
287
+ "model.visual.patch_embed.proj.bias",
288
+ "model.visual.pos_embed.weight",
289
+ "model.visual.merger.norm.weight",
290
+ "model.visual.merger.norm.bias",
291
+ "model.visual.merger.linear_fc1.weight",
292
+ "model.visual.merger.linear_fc1.bias",
293
+ "model.visual.merger.linear_fc2.weight",
294
+ "model.visual.merger.linear_fc2.bias",
295
+ }
296
+ )
297
+ for index in range(int(vision_config.get("depth", 0))):
298
+ prefix = f"model.visual.blocks.{index}"
299
+ required.update(
300
+ {
301
+ f"{prefix}.norm1.weight",
302
+ f"{prefix}.norm1.bias",
303
+ f"{prefix}.norm2.weight",
304
+ f"{prefix}.norm2.bias",
305
+ f"{prefix}.attn.qkv.weight",
306
+ f"{prefix}.attn.qkv.bias",
307
+ f"{prefix}.attn.proj.weight",
308
+ f"{prefix}.attn.proj.bias",
309
+ f"{prefix}.mlp.linear_fc1.weight",
310
+ f"{prefix}.mlp.linear_fc1.bias",
311
+ f"{prefix}.mlp.linear_fc2.weight",
312
+ f"{prefix}.mlp.linear_fc2.bias",
313
+ }
314
+ )
315
+ return required
316
+
317
+
318
+ def _rotate_half(x: torch.Tensor) -> torch.Tensor:
319
+ first, second = x.chunk(2, dim=-1)
320
+ return torch.cat((-second, first), dim=-1)
321
+
322
+
323
+ def _l2norm(x: torch.Tensor, *, eps: float = 1e-6) -> torch.Tensor:
324
+ return x * torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + eps)
325
+
326
+
327
+ def _torch_chunk_gated_delta_rule(
328
+ query: torch.Tensor,
329
+ key: torch.Tensor,
330
+ value: torch.Tensor,
331
+ g: torch.Tensor,
332
+ beta: torch.Tensor,
333
+ *,
334
+ chunk_size: int = 64,
335
+ return_state: bool = False,
336
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
337
+ """Reference chunked gated-delta recurrence used by the transplant path."""
338
+
339
+ initial_dtype = query.dtype
340
+ query = _l2norm(query, eps=1e-6)
341
+ key = _l2norm(key, eps=1e-6)
342
+ query, key, value, beta, g = [
343
+ tensor.transpose(1, 2).contiguous().float() for tensor in (query, key, value, beta, g)
344
+ ]
345
+
346
+ batch_size, num_heads, sequence_length, key_dim = key.shape
347
+ value_dim = value.shape[-1]
348
+ pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
349
+ query = F.pad(query, (0, 0, 0, pad_size))
350
+ key = F.pad(key, (0, 0, 0, pad_size))
351
+ value = F.pad(value, (0, 0, 0, pad_size))
352
+ beta = F.pad(beta, (0, pad_size))
353
+ g = F.pad(g, (0, pad_size))
354
+ total_length = sequence_length + pad_size
355
+ query = query * (key_dim**-0.5)
356
+
357
+ value_beta = value * beta.unsqueeze(-1)
358
+ key_beta = key * beta.unsqueeze(-1)
359
+ query, key, value, key_beta, value_beta = [
360
+ tensor.reshape(tensor.shape[0], tensor.shape[1], -1, chunk_size, tensor.shape[-1])
361
+ for tensor in (query, key, value, key_beta, value_beta)
362
+ ]
363
+ g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
364
+
365
+ upper_including_diagonal = torch.triu(
366
+ torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0
367
+ )
368
+ g = g.cumsum(dim=-1)
369
+ decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
370
+ correction = -((key_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(
371
+ upper_including_diagonal, 0
372
+ )
373
+ for row_index in range(1, chunk_size):
374
+ row = correction[..., row_index, :row_index].clone()
375
+ previous = correction[..., :row_index, :row_index].clone()
376
+ correction[..., row_index, :row_index] = row + (row.unsqueeze(-1) * previous).sum(-2)
377
+ correction = correction + torch.eye(chunk_size, dtype=correction.dtype, device=correction.device)
378
+ value = correction @ value_beta
379
+ cumulative_key = correction @ (key_beta * g.exp().unsqueeze(-1))
380
+
381
+ recurrent = torch.zeros(
382
+ batch_size,
383
+ num_heads,
384
+ key_dim,
385
+ value_dim,
386
+ dtype=value.dtype,
387
+ device=value.device,
388
+ )
389
+ output = torch.zeros_like(value)
390
+ for chunk_index in range(total_length // chunk_size):
391
+ query_chunk = query[:, :, chunk_index]
392
+ key_chunk = key[:, :, chunk_index]
393
+ value_chunk = value[:, :, chunk_index]
394
+ attention = query_chunk @ key_chunk.transpose(-1, -2) * decay_mask[:, :, chunk_index]
395
+ predicted = cumulative_key[:, :, chunk_index] @ recurrent
396
+ corrected_value = value_chunk - predicted
397
+ recurrent_read = (query_chunk * g[:, :, chunk_index, :, None].exp()) @ recurrent
398
+ output[:, :, chunk_index] = recurrent_read + attention @ corrected_value
399
+ recurrent = (
400
+ recurrent * g[:, :, chunk_index, -1, None, None].exp()
401
+ + (
402
+ key_chunk
403
+ * (g[:, :, chunk_index, -1, None] - g[:, :, chunk_index]).exp()[..., None]
404
+ ).transpose(-1, -2)
405
+ @ corrected_value
406
+ )
407
+
408
+ output = output.reshape(output.shape[0], output.shape[1], -1, output.shape[-1])
409
+ output = output[:, :, :sequence_length]
410
+ converted = output.transpose(1, 2).contiguous().to(initial_dtype)
411
+ if return_state:
412
+ return converted, recurrent
413
+ return converted
414
+
415
+
416
+ def _torch_recurrent_gated_delta_rule(
417
+ query: torch.Tensor,
418
+ key: torch.Tensor,
419
+ value: torch.Tensor,
420
+ g: torch.Tensor,
421
+ beta: torch.Tensor,
422
+ state: torch.Tensor,
423
+ ) -> tuple[torch.Tensor, torch.Tensor]:
424
+ """Incremental form numerically aligned with the chunked reference rule with vectorized batched matmuls."""
425
+
426
+ initial_dtype = query.dtype
427
+ query = _l2norm(query, eps=1e-6).float()
428
+ key = _l2norm(key, eps=1e-6).float()
429
+ value = value.float()
430
+ state = state.float()
431
+ key_dim = key.shape[-1]
432
+ scale = float(key_dim**-0.5)
433
+ seq_len = query.shape[1]
434
+
435
+ if seq_len == 1:
436
+ # Fast path for single-token autoregressive decoding (eliminates einsum parsing and dispatch overhead)
437
+ g_exp = g[:, 0, :, None, None].float().exp()
438
+ state = state * g_exp
439
+ k_vec = key[:, 0].unsqueeze(-2)
440
+ prediction = torch.matmul(k_vec, state)
441
+ v_vec = value[:, 0].unsqueeze(-2)
442
+ b_val = beta[:, 0, :, None, None].float()
443
+ delta = b_val * (v_vec - prediction)
444
+ state = state + torch.matmul(k_vec.transpose(-1, -2), delta)
445
+ q_vec = (query[:, 0] * scale).unsqueeze(-2)
446
+ read = torch.matmul(q_vec, state)
447
+ output = read.squeeze(-2).unsqueeze(1).to(initial_dtype)
448
+ return output, state
449
+
450
+ outputs: list[torch.Tensor] = []
451
+ for token_index in range(seq_len):
452
+ g_exp = g[:, token_index, :, None, None].float().exp()
453
+ state = state * g_exp
454
+ k_vec = key[:, token_index].unsqueeze(-2)
455
+ prediction = torch.matmul(k_vec, state)
456
+ v_vec = value[:, token_index].unsqueeze(-2)
457
+ b_val = beta[:, token_index, :, None, None].float()
458
+ delta = b_val * (v_vec - prediction)
459
+ state = state + torch.matmul(k_vec.transpose(-1, -2), delta)
460
+ q_vec = (query[:, token_index] * scale).unsqueeze(-2)
461
+ read = torch.matmul(q_vec, state).squeeze(-2)
462
+ outputs.append(read)
463
+ return torch.stack(outputs, dim=1).to(initial_dtype), state
464
+
465
+
466
+ class DendroExactTransplantCore(SourceBoundModule):
467
+ """Parameterless exact-equivalence operator path inside Dendro Omni.
468
+
469
+ The class reads every tensor from :class:`DendroSourceLayer.named_tensor`.
470
+ Consequently the complete donor and every Dendro extension remain one
471
+ registered parameter tensor.
472
+ """
473
+
474
+ def __init__(self, config: Any, source: DendroSourceLayer) -> None:
475
+ super().__init__(source)
476
+ source_config = dict(config.transplant_source_config or {})
477
+ if not source_config:
478
+ raise ValueError("exact transplant mode requires transplant_source_config")
479
+ self.config = config
480
+ self.source_config = source_config
481
+ self.text_config: Mapping[str, Any] = source_config.get("text_config", source_config)
482
+ self.vision_config: Mapping[str, Any] = source_config.get("vision_config", {})
483
+ self.image_token_id = source_config.get("image_token_id", getattr(config, "image_token_id", None))
484
+ self.video_token_id = source_config.get("video_token_id", getattr(config, "video_token_id", None))
485
+ self._required_prefix = "model.language_model"
486
+ self._bound_weights: dict[tuple[str, torch.dtype | None], torch.Tensor] = {}
487
+ self._bound_source_version: int = -1
488
+
489
+ def _sync_bound_weights(self) -> None:
490
+ current_version = getattr(self.source.source, "_version", 0)
491
+ if self._bound_source_version != current_version:
492
+ self._bound_weights.clear()
493
+ self._bound_source_version = current_version
494
+
495
+ def prebind_weights(self, device: torch.device | str | None = None, dtype: torch.dtype | None = None) -> None:
496
+ """Pre-cache all named tensor slices into an in-memory direct lookup table."""
497
+ self._sync_bound_weights()
498
+ target_device = device or self.source.source.device
499
+ runtime_dtype = dtype or self.source.source.dtype
500
+ for name in self.source._tensor_map:
501
+ t = self.source.named_tensor(name, dtype=runtime_dtype).to(device=target_device)
502
+ self._bound_weights[(name, runtime_dtype)] = t
503
+ self._bound_weights[(name, None)] = t
504
+
505
+ @property
506
+ def _acceleration_backend(self) -> str:
507
+ return str(getattr(self.config, "acceleration_backend", "auto"))
508
+
509
+ def _tensor(self, name: str, *, dtype: torch.dtype | None = None) -> torch.Tensor:
510
+ # Match Hugging Face's ``dtype=...`` loading semantics. The manifest's
511
+ # logical dtype records how each donor tensor was stored, but execution
512
+ # must follow the dtype of the owning single source parameter. Otherwise
513
+ # an FP32 checkpoint silently runs its BF16 donor views in BF16, producing
514
+ # accumulating text drift and much larger vision drift. Callers may still
515
+ # request FP32 explicitly for numerically sensitive state such as A_log.
516
+ self._sync_bound_weights()
517
+ runtime_dtype = self.source.source.dtype if dtype is None else dtype
518
+ key = (name, runtime_dtype)
519
+ cached = self._bound_weights.get(key)
520
+ if cached is not None:
521
+ return cached
522
+ tensor = self.source.named_tensor(name, dtype=runtime_dtype)
523
+ self._bound_weights[key] = tensor
524
+ return tensor
525
+
526
+ def _linear(
527
+ self,
528
+ x: torch.Tensor,
529
+ weight_name: str,
530
+ bias_name: str | None = None,
531
+ ) -> torch.Tensor:
532
+ weight = self._tensor(weight_name, dtype=x.dtype).to(device=x.device)
533
+ bias = None
534
+ if bias_name is not None and self.source.has_named_tensor(bias_name):
535
+ bias = self._tensor(bias_name, dtype=x.dtype).to(device=x.device)
536
+ return F.linear(x, weight, bias)
537
+
538
+ def _qwen_rms_norm(self, x: torch.Tensor, weight_name: str) -> torch.Tensor:
539
+ eps = float(self.text_config.get("rms_norm_eps", 1e-6))
540
+ weight = self._tensor(weight_name, dtype=torch.float32).to(x.device)
541
+ normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps)
542
+ return (normalized * (1.0 + weight.float())).to(x.dtype)
543
+
544
+ def _qwen_gated_rms_norm(
545
+ self,
546
+ x: torch.Tensor,
547
+ gate: torch.Tensor,
548
+ weight_name: str,
549
+ ) -> torch.Tensor:
550
+ eps = float(self.text_config.get("rms_norm_eps", 1e-6))
551
+ weight = self._tensor(weight_name).to(device=x.device, dtype=x.dtype)
552
+ normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps)
553
+ output = weight * normalized.to(x.dtype)
554
+ return (output.float() * F.silu(gate.float())).to(x.dtype)
555
+
556
+ @staticmethod
557
+ def _layer_norm(
558
+ x: torch.Tensor,
559
+ weight: torch.Tensor,
560
+ bias: torch.Tensor,
561
+ eps: float = 1e-6,
562
+ ) -> torch.Tensor:
563
+ return F.layer_norm(
564
+ x,
565
+ (x.shape[-1],),
566
+ weight.to(device=x.device, dtype=x.dtype),
567
+ bias.to(device=x.device, dtype=x.dtype),
568
+ eps,
569
+ )
570
+
571
+ def token_embedding(self, input_ids: torch.Tensor) -> torch.Tensor:
572
+ table = self._tensor(f"{self._required_prefix}.embed_tokens.weight")
573
+ return F.embedding(input_ids, table.to(input_ids.device))
574
+
575
+ def tied_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
576
+ table = self._tensor(f"{self._required_prefix}.embed_tokens.weight")
577
+ return F.linear(hidden_states, table.to(device=hidden_states.device, dtype=hidden_states.dtype))
578
+
579
+ def _text_position_embeddings(
580
+ self,
581
+ hidden_states: torch.Tensor,
582
+ position_ids: torch.Tensor,
583
+ ) -> tuple[torch.Tensor, torch.Tensor]:
584
+ if position_ids.ndim == 2:
585
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
586
+ elif position_ids.ndim == 3 and position_ids.shape[0] == 4:
587
+ position_ids = position_ids[1:]
588
+ if position_ids.ndim != 3 or position_ids.shape[0] != 3:
589
+ raise ValueError("transplant position_ids must be [batch, seq], [3,batch,seq], or [4,batch,seq]")
590
+
591
+ head_dim = int(self.text_config.get("head_dim", 256))
592
+ rope = dict(self.text_config.get("rope_parameters", {}))
593
+ rotary_dim = int(head_dim * float(rope.get("partial_rotary_factor", 1.0)))
594
+ theta = float(rope.get("rope_theta", 10_000.0))
595
+ inv_freq = 1.0 / (
596
+ theta
597
+ ** (
598
+ torch.arange(0, rotary_dim, 2, device=hidden_states.device, dtype=torch.float32)
599
+ / rotary_dim
600
+ )
601
+ )
602
+ frequencies = position_ids.to(device=hidden_states.device, dtype=torch.float32).unsqueeze(-1) * inv_freq
603
+ mixed = frequencies[0].clone()
604
+ sections = list(rope.get("mrope_section", [11, 11, 10]))
605
+ for axis, offset in ((1, 1), (2, 2)):
606
+ length = int(sections[axis]) * 3
607
+ mixed[..., slice(offset, length, 3)] = frequencies[axis, ..., slice(offset, length, 3)]
608
+ embedding = torch.cat((mixed, mixed), dim=-1)
609
+ return embedding.cos().to(hidden_states.dtype), embedding.sin().to(hidden_states.dtype)
610
+
611
+ @staticmethod
612
+ def _apply_text_rope(
613
+ query: torch.Tensor,
614
+ key: torch.Tensor,
615
+ cos: torch.Tensor,
616
+ sin: torch.Tensor,
617
+ ) -> tuple[torch.Tensor, torch.Tensor]:
618
+ cos = cos.unsqueeze(1)
619
+ sin = sin.unsqueeze(1)
620
+ rotary_dim = cos.shape[-1]
621
+ query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
622
+ key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
623
+ query_rotary = query_rotary * cos + _rotate_half(query_rotary) * sin
624
+ key_rotary = key_rotary * cos + _rotate_half(key_rotary) * sin
625
+ return torch.cat((query_rotary, query_pass), -1), torch.cat((key_rotary, key_pass), -1)
626
+
627
+ @staticmethod
628
+ def _repeat_kv(x: torch.Tensor, repetitions: int) -> torch.Tensor:
629
+ if repetitions == 1:
630
+ return x
631
+ batch, heads, length, dim = x.shape
632
+ return x[:, :, None, :, :].expand(batch, heads, repetitions, length, dim).reshape(
633
+ batch, heads * repetitions, length, dim
634
+ )
635
+
636
+ def _full_attention(
637
+ self,
638
+ hidden_states: torch.Tensor,
639
+ layer_index: int,
640
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
641
+ attention_mask: torch.Tensor | None,
642
+ cache: DendroTransplantCache | None = None,
643
+ use_cache: bool = False,
644
+ ) -> torch.Tensor:
645
+ prefix = f"{self._required_prefix}.layers.{layer_index}.self_attn"
646
+ batch, length, _ = hidden_states.shape
647
+ num_heads = int(self.text_config.get("num_attention_heads", 8))
648
+ kv_heads = int(self.text_config.get("num_key_value_heads", num_heads))
649
+ head_dim = int(self.text_config.get("head_dim", hidden_states.shape[-1] // num_heads))
650
+
651
+ mixed_query = self._linear(
652
+ hidden_states, f"{prefix}.q_proj.weight", f"{prefix}.q_proj.bias"
653
+ )
654
+ mixed_query = mixed_query.view(batch, length, num_heads, head_dim * 2)
655
+ query, gate = mixed_query.chunk(2, dim=-1)
656
+ query = self._qwen_rms_norm(query, f"{prefix}.q_norm.weight").transpose(1, 2)
657
+ key = self._linear(
658
+ hidden_states, f"{prefix}.k_proj.weight", f"{prefix}.k_proj.bias"
659
+ ).view(
660
+ batch, length, kv_heads, head_dim
661
+ )
662
+ key = self._qwen_rms_norm(key, f"{prefix}.k_norm.weight").transpose(1, 2)
663
+ value = self._linear(
664
+ hidden_states, f"{prefix}.v_proj.weight", f"{prefix}.v_proj.bias"
665
+ ).view(
666
+ batch, length, kv_heads, head_dim
667
+ ).transpose(1, 2)
668
+ query, key = self._apply_text_rope(query, key, *position_embeddings)
669
+ repetitions = num_heads // kv_heads
670
+
671
+ past_length = 0
672
+ if cache is not None and layer_index in cache.full_attention:
673
+ past_key, past_value = cache.full_attention[layer_index]
674
+ past_key = past_key.to(device=key.device, dtype=key.dtype)
675
+ past_value = past_value.to(device=value.device, dtype=value.dtype)
676
+ past_length = past_key.shape[-2]
677
+ key = torch.cat([past_key, key], dim=-2)
678
+ value = torch.cat([past_value, value], dim=-2)
679
+ if use_cache and cache is not None:
680
+ attention_window = int(getattr(self.config, "transplant_attention_window", 8192))
681
+ cache.full_attention[layer_index] = (
682
+ key[:, :, -attention_window:].detach(),
683
+ value[:, :, -attention_window:].detach(),
684
+ )
685
+
686
+ # Long prefills can use a fused CUDA kernel without an offset mask. Tiny
687
+ # prefills and incremental decode remain faster on the grouped reference
688
+ # operation because kernel selection and transient expansion dominate at
689
+ # those sizes. The persistent GQA cache always remains compact.
690
+ fused_output = None
691
+ # BF16/FP16 fused reductions can move logits enough to alter borderline
692
+ # greedy decisions. Keep reduced-precision inference on the compact,
693
+ # donor-equivalent BMM path; FP32 stays within strict equivalence bounds.
694
+ can_use_exact_inference_fused_prefill = (
695
+ query.dtype == torch.float32
696
+ and attention_mask is None
697
+ and past_length == 0
698
+ and length >= 64
699
+ )
700
+ # Training explicitly opts into fused reduced-precision attention. This
701
+ # is safe for optimization (where small reduction-order differences are
702
+ # expected) and avoids materializing B*H*T*T scores for long examples.
703
+ # Inference remains on the donor-equivalent BF16 reference path. A mask
704
+ # containing padding falls back because combining padding and causality
705
+ # is backend/version sensitive; the audited batch-1 curriculum is fully
706
+ # visible and therefore qualifies.
707
+ training_all_visible = attention_mask is None
708
+ if (
709
+ self.training
710
+ and bool(getattr(self.config, "transplant_training_fused_attention", False))
711
+ and attention_mask is not None
712
+ and attention_mask.ndim == 2
713
+ and tuple(attention_mask.shape) == (batch, key.shape[-2])
714
+ ):
715
+ training_all_visible = bool(
716
+ attention_mask.to(device=query.device, dtype=torch.bool).all().item()
717
+ )
718
+ can_use_training_fused_prefill = (
719
+ self.training
720
+ and bool(getattr(self.config, "transplant_training_fused_attention", False))
721
+ and training_all_visible
722
+ and past_length == 0
723
+ and length >= 64
724
+ )
725
+ if (
726
+ (not self.training and can_use_exact_inference_fused_prefill)
727
+ or can_use_training_fused_prefill
728
+ ):
729
+ fused_output = _run_fused_sdpa(
730
+ query,
731
+ key,
732
+ value,
733
+ is_causal=True,
734
+ enable_gqa=repetitions > 1,
735
+ )
736
+ if fused_output is None and repetitions > 1:
737
+ expanded_key = self._repeat_kv(key, repetitions)
738
+ expanded_value = self._repeat_kv(value, repetitions)
739
+ fused_output = _run_fused_sdpa(
740
+ query,
741
+ expanded_key,
742
+ expanded_value,
743
+ is_causal=True,
744
+ enable_gqa=False,
745
+ )
746
+ if fused_output is not None:
747
+ output = fused_output.transpose(1, 2).contiguous()
748
+ output = output.reshape(batch, length, num_heads * head_dim)
749
+ output = output * torch.sigmoid(gate.reshape(batch, length, num_heads * head_dim))
750
+ return self._linear(output, f"{prefix}.o_proj.weight", f"{prefix}.o_proj.bias")
751
+
752
+ # Preserve the donor's exact reduced-precision arithmetic order. K/V
753
+ # stay compact in the persistent cache, but are expanded transiently for
754
+ # the reference matmuls. A grouped BMM is faster in isolation yet can
755
+ # round borderline BF16 logits differently over long autoregressive runs.
756
+ attention_key = self._repeat_kv(key, repetitions)
757
+ attention_value = self._repeat_kv(value, repetitions)
758
+ scores = torch.matmul(query, attention_key.transpose(-1, -2)) * (head_dim**-0.5)
759
+ query_positions = past_length + torch.arange(length, device=scores.device)
760
+ key_positions = torch.arange(attention_key.shape[-2], device=scores.device)
761
+ causal = key_positions.unsqueeze(0) > query_positions.unsqueeze(1)
762
+ scores = scores.masked_fill(causal[None, None], torch.finfo(scores.dtype).min)
763
+ if attention_mask is not None:
764
+ key_mask = ~attention_mask.to(device=scores.device, dtype=torch.bool)
765
+ scores = scores.masked_fill(
766
+ key_mask[:, None, None, :],
767
+ torch.finfo(scores.dtype).min,
768
+ )
769
+ probabilities = F.softmax(scores.float(), dim=-1).to(query.dtype)
770
+ output = torch.matmul(probabilities, attention_value).transpose(1, 2).contiguous()
771
+ output = output.reshape(batch, length, num_heads * head_dim)
772
+ output = output * torch.sigmoid(gate.reshape(batch, length, num_heads * head_dim))
773
+ return self._linear(output, f"{prefix}.o_proj.weight", f"{prefix}.o_proj.bias")
774
+
775
+ def _linear_attention(
776
+ self,
777
+ hidden_states: torch.Tensor,
778
+ layer_index: int,
779
+ attention_mask: torch.Tensor | None,
780
+ cache: DendroTransplantCache | None = None,
781
+ use_cache: bool = False,
782
+ ) -> torch.Tensor:
783
+ prefix = f"{self._required_prefix}.layers.{layer_index}.linear_attn"
784
+ if attention_mask is not None and hidden_states.shape[0] > 1 and hidden_states.shape[1] > 1:
785
+ hidden_states = hidden_states * attention_mask[:, :, None].to(hidden_states.dtype)
786
+
787
+ batch, length, _ = hidden_states.shape
788
+ key_heads = int(self.text_config.get("linear_num_key_heads", 16))
789
+ value_heads = int(self.text_config.get("linear_num_value_heads", 16))
790
+ key_head_dim = int(self.text_config.get("linear_key_head_dim", 128))
791
+ value_head_dim = int(self.text_config.get("linear_value_head_dim", 128))
792
+ key_dim = key_heads * key_head_dim
793
+ value_dim = value_heads * value_head_dim
794
+
795
+ raw_mixed = self._linear(hidden_states, f"{prefix}.in_proj_qkv.weight").transpose(1, 2)
796
+ convolution_weight = self._tensor(f"{prefix}.conv1d.weight").to(
797
+ device=raw_mixed.device, dtype=raw_mixed.dtype
798
+ )
799
+ if convolution_weight.ndim == 2:
800
+ convolution_weight = convolution_weight.unsqueeze(1)
801
+ kernel = convolution_weight.shape[-1]
802
+ if cache is not None and layer_index in cache.linear_convolution:
803
+ history = cache.linear_convolution[layer_index].to(
804
+ device=raw_mixed.device, dtype=raw_mixed.dtype
805
+ )
806
+ else:
807
+ history = F.pad(raw_mixed[:, :, :0], (max(0, kernel - 1), 0))
808
+ update_state: torch.Tensor | None = None
809
+ accelerated_update = None
810
+ if length == 1 and cache is not None and layer_index in cache.linear_convolution:
811
+ update_state = history
812
+ if update_state.shape[-1] == kernel - 1:
813
+ update_state = F.pad(update_state, (1, 0))
814
+ accelerated_update = fla_causal_conv1d_update(
815
+ raw_mixed.transpose(1, 2).contiguous(),
816
+ update_state.contiguous(),
817
+ convolution_weight.squeeze(1),
818
+ activation="silu",
819
+ configured=self._acceleration_backend,
820
+ )
821
+ if accelerated_update is not None:
822
+ mixed, update_state = accelerated_update
823
+ convolution_input = update_state
824
+ else:
825
+ history_for_convolution = history[..., -(kernel - 1):]
826
+ convolution_input = torch.cat([history_for_convolution, raw_mixed], dim=-1)
827
+ accelerated_mixed = fla_causal_conv1d(
828
+ convolution_input.transpose(1, 2).contiguous(),
829
+ convolution_weight.squeeze(1),
830
+ activation="silu",
831
+ configured=self._acceleration_backend,
832
+ )
833
+ if accelerated_mixed is None:
834
+ mixed = F.silu(F.conv1d(
835
+ convolution_input,
836
+ convolution_weight,
837
+ bias=None,
838
+ padding=0,
839
+ groups=convolution_weight.shape[0],
840
+ )[:, :, :length]).transpose(1, 2)
841
+ else:
842
+ mixed = accelerated_mixed[:, -length:]
843
+ if use_cache and cache is not None:
844
+ history_length = max(0, kernel)
845
+ cache.linear_convolution[layer_index] = (
846
+ convolution_input[:, :, -history_length:].detach()
847
+ if history_length
848
+ else convolution_input[:, :, :0].detach()
849
+ )
850
+ query, key, value = torch.split(mixed, [key_dim, key_dim, value_dim], dim=-1)
851
+ query = query.view(batch, length, key_heads, key_head_dim)
852
+ key = key.view(batch, length, key_heads, key_head_dim)
853
+ value = value.view(batch, length, value_heads, value_head_dim)
854
+
855
+ z = self._linear(hidden_states, f"{prefix}.in_proj_z.weight").view(
856
+ batch, length, value_heads, value_head_dim
857
+ )
858
+ beta = torch.sigmoid(self._linear(hidden_states, f"{prefix}.in_proj_b.weight"))
859
+ a = self._linear(hidden_states, f"{prefix}.in_proj_a.weight")
860
+ a_log = self._tensor(f"{prefix}.A_log", dtype=torch.float32).to(hidden_states.device)
861
+ dt_bias = self._tensor(f"{prefix}.dt_bias", dtype=torch.float32).to(hidden_states.device)
862
+ g = -a_log.float().exp() * F.softplus(a.float() + dt_bias)
863
+
864
+ if value_heads // key_heads > 1:
865
+ repetitions = value_heads // key_heads
866
+ query = query.repeat_interleave(repetitions, dim=2)
867
+ key = key.repeat_interleave(repetitions, dim=2)
868
+ recurrent = None if cache is None else cache.linear_recurrent.get(layer_index)
869
+ if recurrent is None:
870
+ result = fla_chunk_gated_delta_rule(
871
+ query,
872
+ key,
873
+ value,
874
+ g,
875
+ beta,
876
+ return_state=use_cache,
877
+ configured=self._acceleration_backend,
878
+ )
879
+ if result is None:
880
+ result = _torch_chunk_gated_delta_rule(
881
+ query,
882
+ key,
883
+ value,
884
+ g,
885
+ beta,
886
+ return_state=use_cache,
887
+ )
888
+ if use_cache:
889
+ assert isinstance(result, tuple)
890
+ output, recurrent = result
891
+ else:
892
+ assert torch.is_tensor(result)
893
+ output = result
894
+ else:
895
+ accelerated = fla_recurrent_gated_delta_rule(
896
+ query,
897
+ key,
898
+ value,
899
+ g,
900
+ beta,
901
+ recurrent.to(query.device),
902
+ configured=self._acceleration_backend,
903
+ )
904
+ if accelerated is None:
905
+ output, recurrent = _torch_recurrent_gated_delta_rule(
906
+ query,
907
+ key,
908
+ value,
909
+ g,
910
+ beta,
911
+ recurrent.to(query.device),
912
+ )
913
+ else:
914
+ output, recurrent = accelerated
915
+ if use_cache and cache is not None and recurrent is not None:
916
+ cache.linear_recurrent[layer_index] = recurrent.detach()
917
+ output = output.reshape(-1, value_head_dim)
918
+ z = z.reshape(-1, value_head_dim)
919
+ output = self._qwen_gated_rms_norm(output, z, f"{prefix}.norm.weight")
920
+ output = output.reshape(batch, length, value_dim)
921
+ return self._linear(output, f"{prefix}.out_proj.weight")
922
+
923
+ def _mlp(self, hidden_states: torch.Tensor, layer_index: int) -> torch.Tensor:
924
+ prefix = f"{self._required_prefix}.layers.{layer_index}.mlp"
925
+ gate = F.silu(self._linear(hidden_states, f"{prefix}.gate_proj.weight"))
926
+ value = self._linear(hidden_states, f"{prefix}.up_proj.weight")
927
+ return self._linear(gate * value, f"{prefix}.down_proj.weight")
928
+
929
+ def text_forward(
930
+ self,
931
+ *,
932
+ input_ids: torch.Tensor | None = None,
933
+ inputs_embeds: torch.Tensor | None = None,
934
+ attention_mask: torch.Tensor | None = None,
935
+ position_ids: torch.Tensor | None = None,
936
+ cache: DendroTransplantCache | None = None,
937
+ use_cache: bool = False,
938
+ ) -> DendroTransplantOutput:
939
+ if (input_ids is None) == (inputs_embeds is None):
940
+ raise ValueError("provide exactly one of input_ids or inputs_embeds")
941
+ if inputs_embeds is None:
942
+ assert input_ids is not None
943
+ inputs_embeds = self.token_embedding(input_ids)
944
+ batch, length, _ = inputs_embeds.shape
945
+ if use_cache and cache is None:
946
+ cache = DendroTransplantCache()
947
+ supplied_attention_mask = attention_mask
948
+ current_mask = (
949
+ torch.ones(batch, length, device=inputs_embeds.device, dtype=torch.long)
950
+ if attention_mask is None
951
+ else attention_mask.to(device=inputs_embeds.device)
952
+ )
953
+ if current_mask.shape[-1] != length:
954
+ current_mask = current_mask[:, -length:]
955
+ past_mask = None if cache is None else cache.attention_mask
956
+ # ``None`` is the compact representation of an all-visible mask. During
957
+ # evaluation, collapse explicit all-one masks once per forward so every
958
+ # full-attention layer can use maskless fused SDPA. Training stays on the
959
+ # established reference mask path and incurs no synchronization.
960
+ current_all_visible = supplied_attention_mask is None
961
+ if not self.training and not current_all_visible:
962
+ current_all_visible = bool(torch.all(current_mask).item())
963
+ if past_mask is None and current_all_visible:
964
+ attention_mask = None
965
+ else:
966
+ if past_mask is None:
967
+ attention_window = int(getattr(self.config, "transplant_attention_window", 8192))
968
+ past_length = 0 if cache is None else min(int(cache.position), attention_window)
969
+ past_mask = torch.ones(
970
+ batch,
971
+ past_length,
972
+ device=current_mask.device,
973
+ dtype=current_mask.dtype,
974
+ )
975
+ attention_mask = torch.cat([past_mask.to(current_mask.device), current_mask], dim=-1)
976
+ if position_ids is None:
977
+ start = 0 if cache is None else int(cache.position)
978
+ positions = torch.arange(start, start + length, device=inputs_embeds.device, dtype=torch.long)
979
+ position_ids = positions.view(1, 1, -1).expand(4, batch, -1)
980
+ elif position_ids.ndim == 2:
981
+ position_ids = position_ids[None].expand(4, -1, -1)
982
+ position_embeddings = self._text_position_embeddings(inputs_embeds, position_ids)
983
+
984
+ hidden_states = inputs_embeds
985
+ layer_types = list(self.text_config.get("layer_types", []))
986
+ num_layers = int(self.text_config.get("num_hidden_layers", len(layer_types)))
987
+ if not layer_types:
988
+ interval = int(self.text_config.get("full_attention_interval", 4))
989
+ layer_types = [
990
+ "full_attention" if (index + 1) % interval == 0 else "linear_attention"
991
+ for index in range(num_layers)
992
+ ]
993
+ for layer_index in range(num_layers):
994
+ prefix = f"{self._required_prefix}.layers.{layer_index}"
995
+ residual = hidden_states
996
+ normalized = self._qwen_rms_norm(hidden_states, f"{prefix}.input_layernorm.weight")
997
+ if layer_types[layer_index] == "full_attention":
998
+ mixed = self._full_attention(
999
+ normalized,
1000
+ layer_index,
1001
+ position_embeddings,
1002
+ attention_mask,
1003
+ cache=cache,
1004
+ use_cache=use_cache,
1005
+ )
1006
+ else:
1007
+ mixed = self._linear_attention(
1008
+ normalized,
1009
+ layer_index,
1010
+ current_mask,
1011
+ cache=cache,
1012
+ use_cache=use_cache,
1013
+ )
1014
+ hidden_states = residual + mixed
1015
+ residual = hidden_states
1016
+ normalized = self._qwen_rms_norm(hidden_states, f"{prefix}.post_attention_layernorm.weight")
1017
+ hidden_states = residual + self._mlp(normalized, layer_index)
1018
+ hidden_states = self._qwen_rms_norm(hidden_states, f"{self._required_prefix}.norm.weight")
1019
+ if use_cache:
1020
+ assert cache is not None
1021
+ cache.position += length
1022
+ attention_window = int(getattr(self.config, "transplant_attention_window", 8192))
1023
+ cache.attention_mask = (
1024
+ None
1025
+ if attention_mask is None
1026
+ else attention_mask[:, -attention_window:].detach()
1027
+ )
1028
+ return DendroTransplantOutput(
1029
+ last_hidden_state=hidden_states,
1030
+ input_embeddings=inputs_embeds,
1031
+ position_ids=position_ids,
1032
+ cache=cache,
1033
+ )
1034
+
1035
+ # ------------------------------- Vision ---------------------------------
1036
+ @staticmethod
1037
+ def _vision_cu_seqlens(grid_thw: torch.Tensor) -> torch.Tensor:
1038
+ lengths = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0])
1039
+ return F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0), value=0)
1040
+
1041
+ @staticmethod
1042
+ def _vision_position_ids(grid_thw: torch.Tensor, merge_size: int) -> torch.Tensor:
1043
+ positions: list[torch.Tensor] = []
1044
+ for temporal, height, width in grid_thw.tolist():
1045
+ h_grid, w_grid = torch.meshgrid(
1046
+ torch.arange(height, device=grid_thw.device),
1047
+ torch.arange(width, device=grid_thw.device),
1048
+ indexing="ij",
1049
+ )
1050
+ shape = (height // merge_size, merge_size, width // merge_size, merge_size)
1051
+ h_grid = h_grid.reshape(shape).transpose(1, 2).flatten()
1052
+ w_grid = w_grid.reshape(shape).transpose(1, 2).flatten()
1053
+ positions.append(torch.stack((h_grid, w_grid), dim=-1).repeat(temporal, 1))
1054
+ return torch.cat(positions, dim=0)
1055
+
1056
+ @staticmethod
1057
+ def _vision_bilinear_indices(
1058
+ grid_thw: torch.Tensor,
1059
+ side: int,
1060
+ merge_size: int,
1061
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1062
+ index_parts: list[list[torch.Tensor]] = [[] for _ in range(4)]
1063
+ weight_parts: list[list[torch.Tensor]] = [[] for _ in range(4)]
1064
+ device = grid_thw.device
1065
+ for temporal, height, width in grid_thw.tolist():
1066
+ h_grid = torch.linspace(0, side - 1, height, device=device)
1067
+ w_grid = torch.linspace(0, side - 1, width, device=device)
1068
+ h_floor, w_floor = h_grid.int(), w_grid.int()
1069
+ h_ceil = (h_floor + 1).clamp(max=side - 1)
1070
+ w_ceil = (w_floor + 1).clamp(max=side - 1)
1071
+ h_frac, w_frac = h_grid - h_floor, w_grid - w_floor
1072
+ h_floor_offset, h_ceil_offset = h_floor * side, h_ceil * side
1073
+ corners = [
1074
+ (h_floor_offset[:, None] + w_floor[None, :]).flatten(),
1075
+ (h_floor_offset[:, None] + w_ceil[None, :]).flatten(),
1076
+ (h_ceil_offset[:, None] + w_floor[None, :]).flatten(),
1077
+ (h_ceil_offset[:, None] + w_ceil[None, :]).flatten(),
1078
+ ]
1079
+ weights = [
1080
+ ((1 - h_frac)[:, None] * (1 - w_frac)[None, :]).flatten(),
1081
+ ((1 - h_frac)[:, None] * w_frac[None, :]).flatten(),
1082
+ (h_frac[:, None] * (1 - w_frac)[None, :]).flatten(),
1083
+ (h_frac[:, None] * w_frac[None, :]).flatten(),
1084
+ ]
1085
+ h_index = torch.arange(height, device=device).view(height // merge_size, merge_size)
1086
+ w_index = torch.arange(width, device=device).view(width // merge_size, merge_size)
1087
+ reorder = (
1088
+ (h_index[:, :, None, None] * width + w_index[None, None, :, :])
1089
+ .transpose(1, 2)
1090
+ .flatten()
1091
+ .repeat(temporal)
1092
+ )
1093
+ for corner in range(4):
1094
+ index_parts[corner].append(corners[corner][reorder])
1095
+ weight_parts[corner].append(weights[corner][reorder])
1096
+ return (
1097
+ torch.stack([torch.cat(part) for part in index_parts]),
1098
+ torch.stack([torch.cat(part) for part in weight_parts]),
1099
+ )
1100
+
1101
+ @staticmethod
1102
+ def _vision_rope(
1103
+ query: torch.Tensor,
1104
+ key: torch.Tensor,
1105
+ cos: torch.Tensor,
1106
+ sin: torch.Tensor,
1107
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1108
+ query_dtype, key_dtype = query.dtype, key.dtype
1109
+ cos = cos.unsqueeze(-2).float()
1110
+ sin = sin.unsqueeze(-2).float()
1111
+ query = query.float() * cos + _rotate_half(query.float()) * sin
1112
+ key = key.float() * cos + _rotate_half(key.float()) * sin
1113
+ return query.to(query_dtype), key.to(key_dtype)
1114
+
1115
+ def _vision_attention(
1116
+ self,
1117
+ hidden_states: torch.Tensor,
1118
+ layer_index: int,
1119
+ cu_seqlens: torch.Tensor,
1120
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
1121
+ ) -> torch.Tensor:
1122
+ prefix = f"model.visual.blocks.{layer_index}.attn"
1123
+ sequence_length = hidden_states.shape[0]
1124
+ num_heads = int(self.vision_config.get("num_heads", 12))
1125
+ hidden_size = int(self.vision_config.get("hidden_size", hidden_states.shape[-1]))
1126
+ head_dim = hidden_size // num_heads
1127
+ qkv = self._linear(
1128
+ hidden_states,
1129
+ f"{prefix}.qkv.weight",
1130
+ f"{prefix}.qkv.bias",
1131
+ )
1132
+ query, key, value = qkv.reshape(sequence_length, 3, num_heads, head_dim).permute(1, 0, 2, 3).unbind(0)
1133
+ query, key = self._vision_rope(query, key, *position_embeddings)
1134
+ outputs: list[torch.Tensor] = []
1135
+ for start_tensor, end_tensor in zip(cu_seqlens[:-1], cu_seqlens[1:]):
1136
+ start, end = int(start_tensor.item()), int(end_tensor.item())
1137
+ q = query[start:end].transpose(0, 1)
1138
+ k = key[start:end].transpose(0, 1)
1139
+ v = value[start:end].transpose(0, 1)
1140
+ scores = torch.matmul(q, k.transpose(-1, -2)) * (head_dim**-0.5)
1141
+ probabilities = F.softmax(scores.float(), dim=-1).to(q.dtype)
1142
+ outputs.append(torch.matmul(probabilities, v).transpose(0, 1))
1143
+ output = torch.cat(outputs, dim=0).reshape(sequence_length, hidden_size)
1144
+ return self._linear(output, f"{prefix}.proj.weight", f"{prefix}.proj.bias")
1145
+
1146
+ def vision_forward(
1147
+ self,
1148
+ pixel_values: torch.Tensor,
1149
+ grid_thw: torch.Tensor,
1150
+ *,
1151
+ return_intermediates: bool = False,
1152
+ checkpoint_layers: bool = False,
1153
+ ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
1154
+ if not self.vision_config:
1155
+ raise RuntimeError("transplant checkpoint has no vision configuration")
1156
+ hidden_size = int(self.vision_config.get("hidden_size", 768))
1157
+ patch_size = int(self.vision_config.get("patch_size", 16))
1158
+ temporal_patch = int(self.vision_config.get("temporal_patch_size", 2))
1159
+ channels = int(self.vision_config.get("in_channels", 3))
1160
+ merge_size = int(self.vision_config.get("spatial_merge_size", 2))
1161
+ depth = int(self.vision_config.get("depth", 12))
1162
+
1163
+ patch_weight = self._tensor("model.visual.patch_embed.proj.weight")
1164
+ patch_bias = self._tensor("model.visual.patch_embed.proj.bias")
1165
+ patches = pixel_values.view(-1, channels, temporal_patch, patch_size, patch_size)
1166
+ hidden_states = F.conv3d(
1167
+ patches.to(device=patch_weight.device, dtype=patch_weight.dtype),
1168
+ patch_weight,
1169
+ patch_bias,
1170
+ stride=(temporal_patch, patch_size, patch_size),
1171
+ ).view(-1, hidden_size)
1172
+
1173
+ num_positions = int(self.vision_config.get("num_position_embeddings", 2304))
1174
+ side = int(math.sqrt(num_positions))
1175
+ indices, weights = self._vision_bilinear_indices(grid_thw.to(hidden_states.device), side, merge_size)
1176
+ position_table = self._tensor("model.visual.pos_embed.weight").to(hidden_states.device)
1177
+ position_embedding = (position_table[indices] * weights[:, :, None].to(position_table.dtype)).sum(0)
1178
+ hidden_states = hidden_states + position_embedding.to(hidden_states.dtype)
1179
+
1180
+ position_ids = self._vision_position_ids(grid_thw.to(hidden_states.device), merge_size)
1181
+ head_dim = hidden_size // int(self.vision_config.get("num_heads", 12))
1182
+ rotary_dim = head_dim // 2
1183
+ inv_freq = 1.0 / (
1184
+ 10_000.0
1185
+ ** (
1186
+ torch.arange(0, rotary_dim, 2, device=hidden_states.device, dtype=torch.float32)
1187
+ / rotary_dim
1188
+ )
1189
+ )
1190
+ rotary = (position_ids.float().unsqueeze(-1) * inv_freq).flatten(1)
1191
+ rotary = torch.cat((rotary, rotary), dim=-1)
1192
+ position_embeddings = (rotary.cos().to(hidden_states.dtype), rotary.sin().to(hidden_states.dtype))
1193
+ cu_seqlens = self._vision_cu_seqlens(grid_thw.to(hidden_states.device))
1194
+
1195
+ intermediates: list[torch.Tensor] = []
1196
+ capture_layers = {max(0, round((depth - 1) * fraction)) for fraction in (0.25, 0.5, 0.75, 1.0)}
1197
+ for layer_index in range(depth):
1198
+ def vision_block(
1199
+ states: torch.Tensor, current_layer: int = layer_index
1200
+ ) -> torch.Tensor:
1201
+ prefix = f"model.visual.blocks.{current_layer}"
1202
+ norm1 = self._layer_norm(
1203
+ states,
1204
+ self._tensor(f"{prefix}.norm1.weight"),
1205
+ self._tensor(f"{prefix}.norm1.bias"),
1206
+ )
1207
+ states = states + self._vision_attention(
1208
+ norm1, current_layer, cu_seqlens, position_embeddings
1209
+ )
1210
+ norm2 = self._layer_norm(
1211
+ states,
1212
+ self._tensor(f"{prefix}.norm2.weight"),
1213
+ self._tensor(f"{prefix}.norm2.bias"),
1214
+ )
1215
+ mlp = self._linear(
1216
+ norm2,
1217
+ f"{prefix}.mlp.linear_fc1.weight",
1218
+ f"{prefix}.mlp.linear_fc1.bias",
1219
+ )
1220
+ activation_name = str(self.vision_config.get("hidden_act", "gelu")).lower()
1221
+ if activation_name == "gelu_pytorch_tanh":
1222
+ mlp = F.gelu(mlp, approximate="tanh")
1223
+ elif activation_name in {"gelu", "gelu_pytorch"}:
1224
+ mlp = F.gelu(mlp)
1225
+ else:
1226
+ raise ValueError(
1227
+ f"unsupported Qwen3.5 vision activation for exact transplant: {activation_name!r}"
1228
+ )
1229
+ mlp = self._linear(
1230
+ mlp,
1231
+ f"{prefix}.mlp.linear_fc2.weight",
1232
+ f"{prefix}.mlp.linear_fc2.bias",
1233
+ )
1234
+ return states + mlp
1235
+
1236
+ hidden_states = (
1237
+ checkpoint(vision_block, hidden_states, use_reentrant=False)
1238
+ if checkpoint_layers and torch.is_grad_enabled()
1239
+ else vision_block(hidden_states)
1240
+ )
1241
+ if return_intermediates and layer_index in capture_layers:
1242
+ intermediates.append(hidden_states)
1243
+
1244
+ merger_prefix = "model.visual.merger"
1245
+ normalized = self._layer_norm(
1246
+ hidden_states,
1247
+ self._tensor(f"{merger_prefix}.norm.weight"),
1248
+ self._tensor(f"{merger_prefix}.norm.bias"),
1249
+ )
1250
+ merger_hidden = hidden_size * merge_size * merge_size
1251
+ normalized = normalized.view(-1, merger_hidden)
1252
+ merged = self._linear(
1253
+ normalized,
1254
+ f"{merger_prefix}.linear_fc1.weight",
1255
+ f"{merger_prefix}.linear_fc1.bias",
1256
+ )
1257
+ merged = F.gelu(merged)
1258
+ output = self._linear(
1259
+ merged,
1260
+ f"{merger_prefix}.linear_fc2.weight",
1261
+ f"{merger_prefix}.linear_fc2.bias",
1262
+ )
1263
+ if return_intermediates:
1264
+ return output, tuple(intermediates)
1265
+ return output
1266
+
1267
+ def _get_vision_position_ids(
1268
+ self,
1269
+ start_position: int,
1270
+ grid_thw: torch.Tensor,
1271
+ spatial_merge_size: int,
1272
+ device: torch.device,
1273
+ ) -> torch.Tensor:
1274
+ temporal = int(grid_thw[0].item())
1275
+ height = int(grid_thw[1].item()) // spatial_merge_size
1276
+ width = int(grid_thw[2].item()) // spatial_merge_size
1277
+ temporal_positions = torch.arange(temporal, device=device)
1278
+ height_positions = torch.arange(height, device=device) + start_position
1279
+ width_positions = torch.arange(width, device=device) + start_position
1280
+ temporal_grid, height_grid, width_grid = torch.meshgrid(
1281
+ temporal_positions, height_positions, width_positions, indexing="ij"
1282
+ )
1283
+ positions = torch.stack((temporal_grid, height_grid, width_grid), dim=0).reshape(3, -1)
1284
+ positions[0] += start_position
1285
+ return positions
1286
+
1287
+ def get_rope_index(
1288
+ self,
1289
+ input_ids: torch.Tensor,
1290
+ mm_token_type_ids: torch.Tensor,
1291
+ *,
1292
+ image_grid_thw: torch.Tensor | None,
1293
+ video_grid_thw: torch.Tensor | None,
1294
+ attention_mask: torch.Tensor | None,
1295
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1296
+ if video_grid_thw is not None:
1297
+ video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0).clone()
1298
+ video_grid_thw[:, 0] = 1
1299
+ merge_size = int(self.vision_config.get("spatial_merge_size", 2))
1300
+ positions = torch.zeros(
1301
+ 3,
1302
+ input_ids.shape[0],
1303
+ input_ids.shape[1],
1304
+ dtype=input_ids.dtype,
1305
+ device=input_ids.device,
1306
+ )
1307
+ image_iterator = iter(image_grid_thw) if image_grid_thw is not None else None
1308
+ video_iterator = iter(video_grid_thw) if video_grid_thw is not None else None
1309
+ deltas: list[torch.Tensor] = []
1310
+ for batch_index, current_ids in enumerate(input_ids):
1311
+ token_types = mm_token_type_ids[batch_index]
1312
+ active = None
1313
+ if attention_mask is not None:
1314
+ active = attention_mask[batch_index].bool()
1315
+ current_ids = current_ids[active]
1316
+ token_types = token_types[active]
1317
+ groups: list[tuple[int, int, int]] = []
1318
+ for modality, group in itertools.groupby(enumerate(token_types.tolist()), lambda item: item[1]):
1319
+ values = list(group)
1320
+ groups.append((int(modality), values[0][0], values[-1][0] + 1))
1321
+ current_position = 0
1322
+ parts: list[torch.Tensor] = []
1323
+ for modality, start, end in groups:
1324
+ if modality == 0:
1325
+ length = end - start
1326
+ parts.append(
1327
+ torch.arange(length, device=input_ids.device).view(1, -1).expand(3, -1)
1328
+ + current_position
1329
+ )
1330
+ current_position += length
1331
+ else:
1332
+ iterator = image_iterator if modality == 1 else video_iterator
1333
+ if iterator is None:
1334
+ raise ValueError("multimodal token types require matching grid_thw metadata")
1335
+ grid = next(iterator)
1336
+ parts.append(
1337
+ self._get_vision_position_ids(
1338
+ current_position, grid, merge_size, input_ids.device
1339
+ )
1340
+ )
1341
+ current_position += int(max(grid[1], grid[2]).item()) // merge_size
1342
+ current = torch.cat(parts, dim=1)
1343
+ if active is None:
1344
+ positions[:, batch_index] = current
1345
+ else:
1346
+ positions[:, batch_index, active] = current
1347
+ deltas.append(current.max() + 1 - len(current_ids))
1348
+ return positions, torch.stack(deltas).unsqueeze(1)
1349
+
1350
+ def multimodal_forward(
1351
+ self,
1352
+ *,
1353
+ input_ids: torch.Tensor,
1354
+ attention_mask: torch.Tensor | None = None,
1355
+ position_ids: torch.Tensor | None = None,
1356
+ pixel_values: torch.Tensor | None = None,
1357
+ pixel_values_videos: torch.Tensor | None = None,
1358
+ image_grid_thw: torch.Tensor | None = None,
1359
+ video_grid_thw: torch.Tensor | None = None,
1360
+ mm_token_type_ids: torch.Tensor | None = None,
1361
+ cache: DendroTransplantCache | None = None,
1362
+ use_cache: bool = False,
1363
+ ) -> DendroTransplantOutput:
1364
+ embeddings = self.token_embedding(input_ids)
1365
+ if pixel_values is not None:
1366
+ if image_grid_thw is None or self.image_token_id is None:
1367
+ raise ValueError("pixel_values requires image_grid_thw and image_token_id")
1368
+ features = self.vision_forward(pixel_values, image_grid_thw).to(
1369
+ device=embeddings.device, dtype=embeddings.dtype
1370
+ )
1371
+ mask = (input_ids == int(self.image_token_id)).unsqueeze(-1)
1372
+ if int(mask.sum().item()) * embeddings.shape[-1] != features.numel():
1373
+ raise ValueError("image placeholder token count does not match transplanted vision features")
1374
+ embeddings = embeddings.masked_scatter(mask, features)
1375
+ if pixel_values_videos is not None:
1376
+ if video_grid_thw is None or self.video_token_id is None:
1377
+ raise ValueError("pixel_values_videos requires video_grid_thw and video_token_id")
1378
+ features = self.vision_forward(pixel_values_videos, video_grid_thw).to(
1379
+ device=embeddings.device, dtype=embeddings.dtype
1380
+ )
1381
+ mask = (input_ids == int(self.video_token_id)).unsqueeze(-1)
1382
+ if int(mask.sum().item()) * embeddings.shape[-1] != features.numel():
1383
+ raise ValueError("video placeholder token count does not match transplanted vision features")
1384
+ embeddings = embeddings.masked_scatter(mask, features)
1385
+
1386
+ rope_deltas = None
1387
+ if position_ids is None and (image_grid_thw is not None or video_grid_thw is not None):
1388
+ if mm_token_type_ids is None:
1389
+ raise ValueError("multimodal transplant inputs require mm_token_type_ids")
1390
+ position_ids, rope_deltas = self.get_rope_index(
1391
+ input_ids,
1392
+ mm_token_type_ids,
1393
+ image_grid_thw=image_grid_thw,
1394
+ video_grid_thw=video_grid_thw,
1395
+ attention_mask=attention_mask,
1396
+ )
1397
+ result = self.text_forward(
1398
+ inputs_embeds=embeddings,
1399
+ attention_mask=attention_mask,
1400
+ position_ids=position_ids,
1401
+ cache=cache,
1402
+ use_cache=use_cache,
1403
+ )
1404
+ result.rope_deltas = rope_deltas
1405
+ return result
1406
+
1407
+ def validate_required_tensors(self) -> list[str]:
1408
+ """Return missing tensor names needed for the exact path."""
1409
+ required = required_transplant_tensor_names(self.source_config)
1410
+ return sorted(name for name in required if not self.source.has_named_tensor(name))