| --- |
| license: other |
| license_name: modilify-non-commercial-model-license-1.0 |
| license_link: LICENSE |
| library_name: transformers |
| pipeline_tag: text-generation |
| tags: |
| - diffusion |
| - mixture-of-experts |
| - custom-code |
| - modilify-mk2 |
| --- |
| |
| # Modilify Mk2 |
|
|
| Modilify Mk2 is a text diffusion model with learned trajectory memory. Built on the DiffusionGemma 26B-A4B text backbone, it develops a response through repeated revisions of a rolling text canvas. Its memory follows both the current revision process and the content already committed to the response. |
|
|
| This repository provides the complete model weights trained on 78M tokens, tokenizer, and a PyTorch / Transformers runtime. |
|
|
| ## Architecture |
|
|
| Modilify Mk2 works on a 256-token canvas. During each denoising pass, it revises candidate tokens together and uses their confidence and uncertainty to select a prefix to commit. The canvas then advances to continue the response. |
|
|
| Two timescales of GDN2 memory connect these passes. Working memory tracks changes within the canvas, with states for individual positions and the canvas as a whole. Persistent memory incorporates experience from committed tokens and carries it into later parts of the response. The model learns these memory components alongside dense and expert LoRA adapters. |
|
|
| Training supervises response tokens across the valid canvas, with additional objectives for confidence calibration, prefix readiness, and response termination. |
|
|
| | Property | Value | |
| | --- | --- | |
| | Backbone | DiffusionGemma 26B-A4B, text only | |
| | Adaptation | Dense and expert LoRA | |
| | Memory | Dual-timescale GDN2 trajectory memory | |
| | Canvas size | 256 tokens | |
| | Training step | 4400 | |
| | Adaptation tokens | 78,332,410 (prompt + response, excluding padding) | |
| | Weight precision | BF16, with selected memory parameters in FP32 | |
| | Weight files | 16 Safetensors shards, approximately 48.23 GiB | |
| | Framework | PyTorch / Hugging Face Transformers | |
|
|
| ## Quickstart |
|
|
| From the downloaded model directory, use Python 3.12 and install the dependencies: |
|
|
| ```bash |
| python -m pip install -r requirements.txt |
| python inference.py --prompt "Why is the sky blue?" --stream |
| ``` |
|
|
| The CLI selects CUDA, Apple Silicon MPS, or CPU according to availability. Choose a device explicitly with `--device cuda`, `--device mps`, or `--device cpu`. In addition to the weights, allow memory for the prompt cache, trajectory states, and generation workspace. Local inference has been exercised on CPU and MPS. |
|
|
| ## Use with Transformers |
|
|
| The model loads through `AutoModelForCausalLM`. Set `model_path` to the downloaded directory or `Modilify/Modilify-Mk2` on the Hub. |
|
|
| ```python |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| |
| model_path = "./" |
| device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" |
| |
| tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_path, |
| trust_remote_code=True, |
| dtype="auto", |
| device_map=device, |
| attn_implementation="sdpa", |
| ).eval() |
| |
| messages = [{"role": "user", "content": "Explain how rainbows form."}] |
| inputs = tokenizer.apply_chat_template( |
| messages, |
| enable_thinking=True, |
| add_generation_prompt=True, |
| return_tensors="pt", |
| return_dict=True, |
| ).to(model.device) |
| |
| with torch.inference_mode(): |
| output = model.generate( |
| **inputs, |
| max_new_tokens=256, |
| max_denoising_steps=256, |
| seed=42, |
| ) |
| |
| response_ids = output.sequences[0, inputs.input_ids.shape[1]:] |
| print(tokenizer.decode(response_ids, skip_special_tokens=True)) |
| ``` |
|
|
| ### Thinking and conversations |
|
|
| The chat template supports system, user, and assistant turns. Pass the conversation as `messages` and set `enable_thinking=True` to enable reasoning before the answer. The CLI enables thinking by default; `--no-thinking` disables it. |
|
|
| ### Generation controls |
|
|
| `max_new_tokens` sets the response length budget, while `max_denoising_steps` sets the number of canvas revisions. Both default to 256. Increase the budgets for longer responses; reaching either limit can leave a response unfinished. `output.stop_reason` reports how generation ended. |
|
|
| Use `seed` for a single request or `seeds=[...]` for batch requests. For batches, set `tokenizer.padding_side = "left"` and pass the attention mask returned by the tokenizer. The returned `output.sequences` includes the input prompt, and `output.generated_lengths` gives each response's token count. |
|
|
| Generation runs through the included diffusion implementation. The weights include the backbone, adapters, and memory components; `from_pretrained` loads them together. `dtype="auto"` preserves the model's mixed precision settings. |
|
|
| ## Use and limitations |
|
|
| Modilify Mk2 is released for non-commercial research and experimentation with text diffusion, iterative response refinement, and learned generation memory. It accepts text input and produces text, including an optional reasoning channel. |
|
|
| Broad capability and safety benchmark results are not yet available for this checkpoint. Generated responses can contain factual errors, repetition, or incomplete reasoning. Review outputs before using them for consequential decisions. The adaptation dataset is not included in this release. |
|
|
| ## License and credits |
|
|
| The model builds on Google DeepMind's [DiffusionGemma](https://huggingface.co/google/diffusiongemma-26B-A4B-it) and integrates with Hugging Face Transformers. Attribution and upstream licenses are listed in [NOTICE.md](NOTICE.md). |
|
|
| Modilify's contributions are available under the [Modilify Non-Commercial Model License 1.0](LICENSE). Commercial use requires a separate agreement. |
|
|
