Dataset card organized by task

README.md

@@ -4,8 +4,6 @@ base_model: ThinkMorph/ThinkMorph-SSL-Base
 base_model_relation: finetune
 pipeline_tag: any-to-any
 library_name: bagel
-language:
-- en
 datasets:
 - ThinkMorph/Pretext-270K
 - ThinkMorph/Instruct-180K
@@ -13,24 +11,30 @@ datasets:
 
 # ThinkMorph-SSL-Instruct
 
-ThinkMorph-SSL-Instruct is a unified multimodal model that reasons with interleaved text and
-images. On a visual question, it writes reasoning text, generates intermediate images such as a
-zoomed-in crop, a depth map, a bird's-eye view or the view from another position, and then
-gives the answer.
+<p>
+  <a href="https://huggingface.co/datasets/ThinkMorph/Pretext-270K"><img src="https://img.shields.io/badge/Pretext--270K-Dataset-yellow?logo=huggingface&logoColor=yellow" alt="Pretext-270K Dataset"/></a>
+  <a href="https://huggingface.co/datasets/ThinkMorph/Instruct-180K"><img src="https://img.shields.io/badge/Instruct--180K-Dataset-yellow?logo=huggingface&logoColor=yellow" alt="Instruct-180K Dataset"/></a>
+  <a href="https://huggingface.co/ThinkMorph/ThinkMorph-SSL-Base"><img src="https://img.shields.io/badge/ThinkMorph--SSL--Base-Model-orange?logo=huggingface&logoColor=orange" alt="ThinkMorph-SSL-Base Model"/></a>
+</p>
 
-It is built on [BAGEL-7B-MoT](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) in two stages:
+## About
 
-1. **Self-supervised pretext reasoning** on [Pretext-270K](https://huggingface.co/datasets/ThinkMorph/Pretext-270K): the
-   model restores corrupted images and describes them before and after. This stage gives
-   [ThinkMorph-SSL-Base](https://huggingface.co/ThinkMorph/ThinkMorph-SSL-Base).
-2. **Instruction tuning** on [Instruct-180K](https://huggingface.co/datasets/ThinkMorph/Instruct-180K): question-driven
-   traces whose intermediate images come from 17 visual operations.
+ThinkMorph-SSL-Instruct reasons with interleaved text and images. On a visual question it thinks
+in text, generates intermediate images when they help (a zoomed-in crop, a depth map, a
+bird's-eye view, the view from another position), and then answers.
 
-This repository holds the EMA weights at step 5,500 of the second stage.
+It is trained from [BAGEL-7B-MoT](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) in two stages:
 
-## Prompt format
+| Stage | Data | Model |
+| --- | --- | --- |
+| 1. Self-supervised pretext reasoning | [Pretext-270K](https://huggingface.co/datasets/ThinkMorph/Pretext-270K) | [ThinkMorph-SSL-Base](https://huggingface.co/ThinkMorph/ThinkMorph-SSL-Base) |
+| 2. Instruction tuning | [Instruct-180K](https://huggingface.co/datasets/ThinkMorph/Instruct-180K) | ThinkMorph-SSL-Instruct (this model: EMA weights, step 5,500) |
+
+## Usage
 
-The model was trained with the system prompt below. Use it at inference:
+The checkpoint has the same layout as BAGEL-7B-MoT and loads with the
+[BAGEL code](https://github.com/ByteDance-Seed/Bagel). Keep the VAE and tokenizer files shipped
+here. Prompt the model with the system prompt it was trained with:
 
 ```text
 Use <think>...</think> for intermediate reasoning or planning.
@@ -49,54 +53,10 @@ Example:
   <answer>final answer</answer>
 ```
 
-When the model writes `<gen_image>[Image size: (height, width)]`, the inference loop generates
-an image with the model's image branch, adds it to the context and continues decoding.
-
-## Loading
-
-The repository has the same layout as BAGEL-7B-MoT, so the official BAGEL inference code
-(<https://github.com/ByteDance-Seed/Bagel>) loads it: point `model_path` at a local copy.
-
-```bash
-hf download ThinkMorph/ThinkMorph-SSL-Instruct --local-dir ThinkMorph-SSL-Instruct
-```
-
-Keep the VAE and the tokenizer files shipped here. The VAE is the one used in training, and
-the tokenizer has the special tokens the model was trained with.
-
-## Files
-
-| File | Content |
-| --- | --- |
-| `ema.safetensors` | the model weights: bfloat16, 14.7B parameters (29.4 GB) |
-| `ae.safetensors` | the VAE, frozen during training |
-| `config.json`, `llm_config.json`, `vit_config.json`, `generation_config.json` | model configuration |
-| `vocab.json`, `merges.txt`, `added_tokens.json`, `special_tokens_map.json`, `tokenizer_config.json` | the tokenizer, with the special tokens used in training |
-
-## Training
-
-| | Stage 1: pretext reasoning | Stage 2: instruction tuning |
-| --- | --- | --- |
-| Starts from | BAGEL-7B-MoT | ThinkMorph-SSL-Base |
-| Data | Pretext-270K, 270,896 training traces | Instruct-180K, 180,490 traces |
-| GPUs | 32 | 48 |
-| Learning rate, warmup | 2.83e-5, 1,000 steps | 2e-5, 1,000 steps |
-| Sequence packing | up to 16,384 tokens | up to 20,480 tokens |
-| Released weights | EMA at step 5,500 | EMA at step 5,500 (this repository) |
-
-Both stages train on next-token cross-entropy for text and flow matching for generated images,
-with weight 1 each. The VAE is frozen; the language model and the vision encoder are trained.
-EMA weights are kept in fp32 with decay 0.9998 from step 1,000.
-
-## Evaluation
+## Benchmarks
 
 Results will be added with the paper.
 
-## License
-
-The weights are released under Apache-2.0, following BAGEL-7B-MoT. The training data keep
-the terms of their image sources; see the dataset cards.
-
 ## Citation
 
 The citation will be added with the paper.