Dataset card organized by task

README.md

@@ -4,29 +4,33 @@ base_model: ByteDance-Seed/BAGEL-7B-MoT
 base_model_relation: finetune
 pipeline_tag: any-to-any
 library_name: bagel
-language:
-- en
 datasets:
 - ThinkMorph/Pretext-270K
 ---
 
 # ThinkMorph-SSL-Base
 
-ThinkMorph-SSL-Base is [BAGEL-7B-MoT](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) after self-supervised pretext
-reasoning on [Pretext-270K](https://huggingface.co/datasets/ThinkMorph/Pretext-270K). Each training trace starts from
-a corrupted image: grayscale, noisy, low-resolution, masked, warped, rotated or flipped, or cut
-into shuffled tiles. The model describes the input, generates the restored image, and describes
-the result.
+<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/ThinkMorph/ThinkMorph-SSL-Instruct"><img src="https://img.shields.io/badge/ThinkMorph--SSL--Instruct-Model-orange?logo=huggingface&logoColor=orange" alt="ThinkMorph-SSL-Instruct Model"/></a>
+</p>
+
+## About
 
-This repository holds the EMA weights at step 5,500.
-[ThinkMorph-SSL-Instruct](https://huggingface.co/ThinkMorph/ThinkMorph-SSL-Instruct) is trained from these weights.
+ThinkMorph-SSL-Base is [BAGEL-7B-MoT](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) after self-supervised pretext
+reasoning on [Pretext-270K](https://huggingface.co/datasets/ThinkMorph/Pretext-270K). In each training trace the model
+sees a corrupted image (grayscale, noisy, low-resolution, masked, warped, rotated or flipped, or
+cut into shuffled tiles), describes it, generates the restored image, and describes the result.
 
-The model has seen no questions or answers in training. Use it as a starting point for further
-training. For visual question answering, use ThinkMorph-SSL-Instruct.
+This repository holds the EMA weights at step 5,500, the starting point of
+[ThinkMorph-SSL-Instruct](https://huggingface.co/ThinkMorph/ThinkMorph-SSL-Instruct). The model has seen no
+questions or answers; for visual question answering, use ThinkMorph-SSL-Instruct.
 
-## Prompt format
+## 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.
@@ -45,54 +49,6 @@ 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-Base --local-dir ThinkMorph-SSL-Base
-```
-
-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 |
-
-All files except `ema.safetensors` are identical to those of ThinkMorph-SSL-Instruct. The EMA
-state of the training checkpoint was merged and cast to bfloat16 in the same layout (1,451
-tensors). The unused last SigLIP layer is left out, as in BAGEL.
-
-## Training
-
-| | Stage 1: pretext reasoning |
-| --- | --- |
-| Starts from | BAGEL-7B-MoT |
-| Data | Pretext-270K, 270,896 training traces |
-| GPUs | 32 |
-| Learning rate, warmup | 2.83e-5, 1,000 steps |
-| Sequence packing | up to 16,384 tokens |
-| Released weights | EMA at step 5,500 |
-
-Training uses 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.
-
-## 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 card.
-
 ## Citation
 
 The citation will be added with the paper.