Image-Text-to-Text
Transformers
Safetensors
English
dendro_omni
text-generation
phillnet
phillnet-mini
dendro
visual-question-answering
multimodal
adaptive-reasoning
code-generation
long-context
custom-code
text-vision-only
conversational
custom_code
Instructions to use ayjays132/Phillnet-Mini-Max with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use ayjays132/Phillnet-Mini-Max with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("image-text-to-text", model="ayjays132/Phillnet-Mini-Max", trust_remote_code=True) messages = [ { "role": "user", "content": [ {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"}, {"type": "text", "text": "What animal is on the candy?"} ] }, ] pipe(text=messages)# Load model directly from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("ayjays132/Phillnet-Mini-Max", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use ayjays132/Phillnet-Mini-Max with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "ayjays132/Phillnet-Mini-Max" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "ayjays132/Phillnet-Mini-Max", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker
docker model run hf.co/ayjays132/Phillnet-Mini-Max
- SGLang
How to use ayjays132/Phillnet-Mini-Max with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "ayjays132/Phillnet-Mini-Max" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "ayjays132/Phillnet-Mini-Max", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "ayjays132/Phillnet-Mini-Max" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "ayjays132/Phillnet-Mini-Max", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }' - Docker Model Runner
How to use ayjays132/Phillnet-Mini-Max with Docker Model Runner:
docker model run hf.co/ayjays132/Phillnet-Mini-Max
Release Phillnet Mini Text-Vision v1.1.0
Browse filesText-and-still-image model release with adaptive reasoning defaults, complete model card, examples, integrity manifest, and production deployment materials.
- .dockerignore +15 -0
- .env.example +14 -0
- .gitattributes +7 -34
- .gitignore +29 -0
- ADAPTIVE_DEFAULTS_REPORT.md +79 -0
- COHERENCE_REPAIR_REPORT.md +71 -0
- Dockerfile +22 -0
- HF_UPLOAD.md +61 -0
- IMPLEMENTATION_REPORT.md +240 -0
- PRODUCTION.md +108 -0
- README.md +285 -0
- RELEASE_MANIFEST.json +59 -0
- RELEASE_VALIDATION.md +27 -0
- __init__.py +33 -0
- chat_template.jinja +154 -0
- config.json +0 -0
- docker-compose.yml +23 -0
- examples/direct-horizon-tasks.png +3 -0
- examples/max-orbit-notes.png +0 -0
- examples/medium-solara-dashboard.png +3 -0
- modalities.py +407 -0
- modeling_dendro_omni.py +1929 -0
- preprocessor_config.json +15 -0
- processing_dendro_omni.py +267 -0
- requirements-server.txt +3 -0
- server.py +226 -0
- tokenization_dendro_omni.py +334 -0
- tokenizer_config.json +32 -0
- transplant.py +1410 -0
.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 |
-
|
| 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 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 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 |
+
| [](examples/direct-horizon-tasks.html) | [](examples/medium-solara-dashboard.html) | [](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
|
examples/max-orbit-notes.png
ADDED
|
examples/medium-solara-dashboard.png
ADDED
|
Git LFS Details
|
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))
|