Add self-contained DiffRetriever (trust_remote_code: code + config + adapter/)

README.md

@@ -0,0 +1,206 @@
+---
+base_model: Dream-org/Dream-v0-Instruct-7B
+library_name: peft
+tags:
+- base_model:adapter:Dream-org/Dream-v0-Instruct-7B
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.18.1

adapter/adapter_config.json

@@ -0,0 +1,46 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": null,
+  "base_model_name_or_path": "Dream-org/Dream-v0-Instruct-7B",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 64,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "peft_version": "0.18.1",
+  "qalora_group_size": 16,
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "v_proj",
+    "q_proj",
+    "k_proj",
+    "o_proj",
+    "up_proj",
+    "gate_proj",
+    "down_proj"
+  ],
+  "target_parameters": null,
+  "task_type": "FEATURE_EXTRACTION",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

adapter/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2963a476388a17cf51e47c1141d00ff7c9fb8bd575459862d632cf9749becfac
+size 80792880

added_tokens.json

@@ -0,0 +1,26 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|beginoftext|>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|mask|>": 151666,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

backbone_adapters.py

@@ -0,0 +1,329 @@
+"""
+Backbone adapters for diffusion retriever training.
+
+Each adapter encapsulates ALL model-specific behavior in one place:
+  - How to load the backbone (AutoModel vs AutoModelForCausalLM)
+  - PEFT/LoRA configuration (target modules, task type)
+  - Attention mask format (2D vs 4D)
+  - Hidden state extraction (forward hook on output projection)
+  - Mask token ID (verified from HuggingFace tokenizer configs)
+  - Gradient checkpointing support
+
+The TrainableDiffusionRetriever delegates to an adapter and has ZERO
+model-specific branches.  Adding a new model = adding one adapter class.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Tuple
+import logging
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Base adapter
+# ---------------------------------------------------------------------------
+
+class BackboneAdapter(ABC):
+    """Abstract interface for model-specific backbone behavior."""
+
+    model_type: str          # e.g. 'dream', 'llada1', 'llada2'
+    mask_token_id: int       # verified from HuggingFace model cards
+    hub_model_name: str      # HuggingFace model ID for fallback loading
+
+    def __init__(self):
+        self.flash_attn: bool = False   # set by load_backbone
+
+    # -- loading --------------------------------------------------------
+
+    def load_backbone(self, source: str, device_map=None) -> nn.Module:
+        """Load backbone from HuggingFace model name or local directory.
+
+        Tries flash attention variants in order, falls back to eager.
+        Sets ``self.flash_attn`` as a side effect.
+        """
+        common_kw = dict(trust_remote_code=True, torch_dtype=torch.bfloat16)
+        if device_map is not None:
+            common_kw['device_map'] = device_map
+
+        for attn_impl in self._flash_attn_impls:
+            try:
+                bb = self._auto_class().from_pretrained(
+                    source, attn_implementation=attn_impl, **common_kw)
+                self.flash_attn = True
+                logger.info(f"{self.model_type}: {attn_impl} enabled")
+                return bb
+            except (ValueError, ImportError):
+                pass
+
+        self.flash_attn = False
+        return self._auto_class().from_pretrained(source, **common_kw)
+
+    @staticmethod
+    @abstractmethod
+    def _auto_class():
+        """Return the AutoModel class to use (AutoModel or AutoModelForCausalLM)."""
+
+    _flash_attn_impls: Tuple[str, ...] = ('flash_attention_2',)
+
+    # -- PEFT / LoRA ----------------------------------------------------
+
+    @abstractmethod
+    def get_lora_config(self, lora_rank: int, lora_alpha: int,
+                        lora_dropout: float = 0.0):
+        """Return a ``peft.LoraConfig`` appropriate for this backbone."""
+
+    # -- attention mask --------------------------------------------------
+
+    @abstractmethod
+    def needs_4d_mask(self) -> bool:
+        """Whether the backbone expects a 4D ``[B,1,L,L]`` attention mask.
+
+        If False, the backbone handles bidirectional attention internally
+        and expects a standard 2D ``[B,L]`` padding mask.
+        """
+
+    # -- hidden state extraction -----------------------------------------
+
+    def register_hidden_hook(self, backbone: nn.Module,
+                             ref_dict: Dict[str, torch.Tensor]) -> bool:
+        """Register a forward hook on the output projection to capture the
+        last hidden state without ``output_hidden_states=True``.
+
+        Returns True if a hook was registered, False otherwise (in which
+        case the caller should fall back to ``output_hidden_states``).
+        """
+        return False  # default: no hook, use output_hidden_states
+
+    # -- gradient checkpointing ------------------------------------------
+
+    def enable_gradient_checkpointing(self, backbone: nn.Module, **kwargs):
+        """Enable gradient checkpointing.  Override for models that don't
+        support it."""
+        backbone.gradient_checkpointing_enable(**kwargs)
+        logger.info("Gradient checkpointing enabled")
+
+
+# ---------------------------------------------------------------------------
+# Hook helpers (shared across adapters)
+# ---------------------------------------------------------------------------
+
+def _is_linear(mod: nn.Module) -> bool:
+    """Check if a module is a Linear layer (plain or LoRA-wrapped)."""
+    if isinstance(mod, nn.Linear):
+        return True
+    # PEFT LoRA wraps nn.Linear in peft.tuners.lora.layer.Linear which is
+    # NOT a subclass of nn.Linear, but has a base_layer that is.
+    if hasattr(mod, 'base_layer') and isinstance(mod.base_layer, nn.Linear):
+        return True
+    return False
+
+
+def _hook_on_module(backbone: nn.Module, ref_dict: Dict,
+                    target_name: str, skip_if_contains: Optional[str] = None,
+                    adapter_name: str = '') -> bool:
+    """Register a forward hook on the first Linear (or LoRA-wrapped Linear)
+    whose leaf name matches *target_name*.  Optionally skip modules whose
+    full path contains *skip_if_contains* (e.g. 'blocks' to skip per-layer
+    ff_out).
+    """
+    for name, mod in backbone.named_modules():
+        leaf = name.split('.')[-1]
+        if leaf == target_name and _is_linear(mod):
+            if skip_if_contains and skip_if_contains in name:
+                continue
+            mod.register_forward_hook(
+                lambda m, inp, out, r=ref_dict: r.update({'h': inp[0]})
+            )
+            logger.info(f"{adapter_name}: hook on '{name}'")
+            return True
+    return False
+
+
+# ---------------------------------------------------------------------------
+# Dream
+# ---------------------------------------------------------------------------
+
+class DreamAdapter(BackboneAdapter):
+    model_type = 'dream'
+    mask_token_id = 151666          # <|mask|>  — Dream-org/Dream-v0-Instruct-7B
+    hub_model_name = 'Dream-org/Dream-v0-Instruct-7B'
+
+    @staticmethod
+    def _auto_class():
+        from transformers import AutoModel
+        return AutoModel
+
+    def get_lora_config(self, lora_rank, lora_alpha, lora_dropout=0.0):
+        from peft import LoraConfig, TaskType
+        # Dream is loaded via AutoModel (not AutoModelForCausalLM), so
+        # FEATURE_EXTRACTION avoids PeftModelForCausalLM which would
+        # require prepare_inputs_for_generation (Dream lacks this).
+        return LoraConfig(
+            r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                            "gate_proj", "up_proj", "down_proj"],
+            task_type=TaskType.FEATURE_EXTRACTION,
+            bias="none",
+        )
+
+    def needs_4d_mask(self) -> bool:
+        # Standard HF Qwen2 attention — needs 4D to enforce bidirectional.
+        return True
+
+    def register_hidden_hook(self, backbone, ref_dict):
+        return _hook_on_module(backbone, ref_dict, 'lm_head',
+                               adapter_name='dream')
+
+
+# ---------------------------------------------------------------------------
+# LLaDA v1 (GSAI-ML/LLaDA-8B-Instruct)
+# ---------------------------------------------------------------------------
+
+class LLaDA1Adapter(BackboneAdapter):
+    model_type = 'llada1'
+    mask_token_id = 126336          # <|mdm_mask|>  — GSAI-ML/LLaDA-8B-Instruct
+    hub_model_name = 'GSAI-ML/LLaDA-8B-Instruct'
+    _flash_attn_impls = ('flash_attention_3', 'flash_attention_2')
+
+    @staticmethod
+    def _auto_class():
+        from transformers import AutoModelForCausalLM
+        return AutoModelForCausalLM
+
+    def get_lora_config(self, lora_rank, lora_alpha, lora_dropout=0.0):
+        from peft import LoraConfig, TaskType
+        # LLaDA1 custom arch uses different names than standard LLaMA:
+        #   attn_out (not o_proj), ff_proj (not gate_proj), ff_out (not down_proj)
+        return LoraConfig(
+            r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
+            target_modules=["q_proj", "k_proj", "v_proj", "attn_out",
+                            "ff_proj", "up_proj", "ff_out"],
+            task_type=TaskType.CAUSAL_LM,
+            bias="none",
+        )
+
+    def needs_4d_mask(self) -> bool:
+        # Custom code handles bidirectional attention internally — always 2D.
+        return False
+
+    def register_hidden_hook(self, backbone, ref_dict):
+        # LLaDA1 has ff_out per-block (blocks.X.ff_out = FFN output) AND at
+        # the model level (transformer.ff_out = output projection to vocab).
+        # We need the model-level one; skip per-block ones.
+        return _hook_on_module(backbone, ref_dict, 'ff_out',
+                               skip_if_contains='blocks',
+                               adapter_name='llada1')
+
+    def enable_gradient_checkpointing(self, backbone, **kwargs):
+        # LLaDA1's LLaDAModelLM doesn't support HF gradient_checkpointing_enable.
+        # Manually wrap each transformer block with torch checkpoint.
+        from torch.utils.checkpoint import checkpoint as ckpt_fn
+
+        # Find the blocks ModuleList through the PEFT wrapper
+        blocks = None
+        for name, mod in backbone.named_modules():
+            if name.endswith('.blocks') and isinstance(mod, nn.ModuleList):
+                blocks = mod
+                break
+
+        if blocks is None:
+            logger.warning("LLaDA1: couldn't find transformer blocks — "
+                           "skipping gradient checkpointing")
+            return
+
+        for block in blocks:
+            orig_forward = block.forward
+
+            def _make_ckpt(fwd):
+                def _ckpt_forward(*args, **kwargs):
+                    if not torch.is_grad_enabled():
+                        return fwd(*args, **kwargs)
+                    return ckpt_fn(fwd, *args, use_reentrant=False, **kwargs)
+                return _ckpt_forward
+
+            block.forward = _make_ckpt(orig_forward)
+
+        logger.info(f"LLaDA1: manual gradient checkpointing on {len(blocks)} blocks")
+
+
+# ---------------------------------------------------------------------------
+# LLaDA v1.5 (GSAI-ML/LLaDA-1.5)  — same architecture as v1
+# ---------------------------------------------------------------------------
+
+class LLaDA15Adapter(LLaDA1Adapter):
+    model_type = 'llada15'
+    mask_token_id = 126336          # <|mdm_mask|>  — same tokenizer as v1
+    hub_model_name = 'GSAI-ML/LLaDA-1.5'
+
+    def register_hidden_hook(self, backbone, ref_dict):
+        return _hook_on_module(backbone, ref_dict, 'ff_out',
+                               skip_if_contains='blocks',
+                               adapter_name='llada15')
+
+
+
+# ---------------------------------------------------------------------------
+# LLaDA v2 (inclusionAI/LLaDA2.0-mini)
+# ---------------------------------------------------------------------------
+
+class LLaDA2Adapter(BackboneAdapter):
+    model_type = 'llada2'
+    mask_token_id = 156895          # <|mask|>  — inclusionAI/LLaDA2.0-mini
+    hub_model_name = 'inclusionAI/LLaDA2.0-mini'
+    _flash_attn_impls = ('flash_attention_3', 'flash_attention_2')
+
+    @staticmethod
+    def _auto_class():
+        from transformers import AutoModelForCausalLM
+        return AutoModelForCausalLM
+
+    def get_lora_config(self, lora_rank, lora_alpha, lora_dropout=0.0):
+        from peft import LoraConfig, TaskType
+        # LLaDA2 uses fused QKV ("query_key_value") and "dense" for attn
+        # output.  Skip MoE expert FFN layers to avoid multiplying params.
+        return LoraConfig(
+            r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
+            target_modules=[
+                "query_key_value", "dense",
+                "mlp.shared_experts.gate_proj",
+                "mlp.shared_experts.up_proj",
+                "mlp.shared_experts.down_proj",
+            ],
+            task_type=TaskType.CAUSAL_LM,
+            bias="none",
+        )
+
+    def needs_4d_mask(self) -> bool:
+        # Standard HF causal model — needs 4D to override causal attention,
+        # unless flash attention handles masking itself.
+        return not self.flash_attn
+
+    def register_hidden_hook(self, backbone, ref_dict):
+        return _hook_on_module(backbone, ref_dict, 'lm_head',
+                               adapter_name='llada2')
+
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+
+ADAPTER_REGISTRY: Dict[str, type] = {
+    'dream':  DreamAdapter,
+    'llada1': LLaDA1Adapter,
+    'llada15': LLaDA15Adapter,
+    'llada2': LLaDA2Adapter,
+}
+
+
+def get_adapter(model_type: str) -> BackboneAdapter:
+    """Create a BackboneAdapter for the given model_type."""
+    cls = ADAPTER_REGISTRY.get(model_type)
+    if cls is None:
+        raise ValueError(
+            f"Unknown model_type: {model_type!r}. "
+            f"Available: {sorted(ADAPTER_REGISTRY.keys())}")
+    return cls()

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,15 @@
+{
+  "model_type": "diffretriever",
+  "architectures": [
+    "DiffRetrieverModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_diffretriever.DiffRetrieverConfig",
+    "AutoModel": "modeling_diffretriever.DiffRetrieverModel"
+  },
+  "base_model": "Dream-org/Dream-v0-Instruct-7B",
+  "backbone_type": "dream",
+  "mode": "multi",
+  "k_q": 4,
+  "k_p": 16
+}

configuration_diffretriever.py

@@ -0,0 +1,36 @@
+"""HuggingFace config for DiffRetriever `trust_remote_code` loading.
+
+A thin `PretrainedConfig` so that
+
+    AutoModel.from_pretrained("ielabgroup/diffretriever-...", trust_remote_code=True)
+
+can route to `DiffRetrieverModel` via the repo's `config.json` `auto_map`.
+
+The real retrieval configuration (prompt token ids, K_q/K_p, temperature,
+sparse weight, ...) lives in `retriever_config.json` and is read by
+`TrainableDiffusionRetriever.load()`. The fields here are informational only
+(they show up in the Hub config viewer) and are not required for loading.
+
+This file is shipped *inside each model repo* — keep it dependency-light.
+"""
+from transformers import PretrainedConfig
+
+
+class DiffRetrieverConfig(PretrainedConfig):
+    model_type = "diffretriever"
+
+    def __init__(
+        self,
+        base_model: str | None = None,
+        backbone_type: str | None = None,
+        mode: str = "single",
+        k_q: int = 1,
+        k_p: int = 1,
+        **kwargs,
+    ):
+        self.base_model = base_model      # e.g. "Dream-org/Dream-v0-Instruct-7B"
+        self.backbone_type = backbone_type  # e.g. "dream" / "llada"
+        self.mode = mode                  # "single" | "multi"
+        self.k_q = k_q
+        self.k_p = k_p
+        super().__init__(**kwargs)

modeling_diffretriever.py

@@ -0,0 +1,104 @@
+"""HuggingFace `trust_remote_code` wrapper for DiffRetriever.
+
+Lets a released checkpoint load with a single call:
+
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained(
+        "ielabgroup/diffretriever-dream-7b-single", trust_remote_code=True)
+    model.eval()
+    ids, mask = model.tokenize(["a query"], is_query=True)
+    out = model.encode(ids, mask, is_query=True)   # {'repr_hidden', ...}
+
+It wraps `TrainableDiffusionRetriever` (shipped in the same repo) and exposes
+`.tokenize()` / `.encode()` / `.backbone`. The base diffusion backbone
+(Dream / LLaDA / ...) is pulled from its own Hub repo at load time; this repo
+carries only the LoRA adapter + tokenizer + retriever_config.json + this code.
+
+Shipped inside each model repo, so keep the import surface minimal.
+"""
+from __future__ import annotations
+
+import os
+import shutil
+import tempfile
+
+import torch  # noqa: F401  (used by the wrapped retriever; keep import explicit)
+from transformers import PreTrainedModel
+
+from .configuration_diffretriever import DiffRetrieverConfig
+from .diffretriever_trainable import TrainableDiffusionRetriever
+
+
+class DiffRetrieverModel(PreTrainedModel):
+    config_class = DiffRetrieverConfig
+
+    def __init__(self, config: DiffRetrieverConfig, retriever=None):
+        super().__init__(config)
+        # Registered as a submodule so .to()/.eval()/.parameters() recurse into it.
+        self.retriever = retriever
+
+    # ── Loading ──────────────────────────────────────────────────────────────
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        # We do NOT call super().from_pretrained: the checkpoint is a LoRA
+        # adapter that must be attached on top of a base backbone pulled from a
+        # different Hub repo, which TrainableDiffusionRetriever.load() handles.
+        kwargs.pop("trust_remote_code", None)
+        config = kwargs.pop("config", None)
+
+        path = str(pretrained_model_name_or_path)
+        if not os.path.isdir(path):
+            from huggingface_hub import snapshot_download
+            dl = {k: kwargs[k] for k in
+                  ("cache_dir", "revision", "token", "local_files_only", "proxies")
+                  if k in kwargs and kwargs[k] is not None}
+            path = snapshot_download(path, **dl)
+
+        # Build a temp view of the snapshot for TrainableDiffusionRetriever.load:
+        #   - DROP config.json: load() detects a released adapter repo only when
+        #     adapter_config.json is present AND config.json is absent; this repo
+        #     ships a config.json for the auto_map.
+        #   - FLATTEN the adapter/ subdir into the root: the LoRA adapter is
+        #     stored under adapter/ (not the repo root) on purpose, so that
+        #     transformers' PEFT auto-loader does NOT hijack AutoModel and load
+        #     the base model directly instead of this wrapper. load() still needs
+        #     adapter_config.json at the top level, so we link it up here.
+        tmp = tempfile.mkdtemp(prefix="diffretriever_")
+
+        def _link(src, dst):
+            if os.path.isfile(src) and not os.path.exists(dst):
+                os.symlink(os.path.abspath(src), dst)
+
+        try:
+            for fn in os.listdir(path):
+                if fn == "config.json":
+                    continue
+                _link(os.path.join(path, fn), os.path.join(tmp, fn))
+            adapter_dir = os.path.join(path, "adapter")
+            if os.path.isdir(adapter_dir):
+                for fn in os.listdir(adapter_dir):
+                    _link(os.path.join(adapter_dir, fn), os.path.join(tmp, fn))
+            retriever = TrainableDiffusionRetriever.load(tmp)
+        finally:
+            shutil.rmtree(tmp, ignore_errors=True)
+
+        if config is None:
+            config = DiffRetrieverConfig(
+                base_model=getattr(retriever, "model_name", None),
+                backbone_type=getattr(retriever, "model_type", None),
+            )
+        return cls(config, retriever=retriever)
+
+    # ── Retrieval API (delegates to the wrapped retriever) ─────────────────────
+    @property
+    def backbone(self):
+        return self.retriever.backbone
+
+    def tokenize(self, *args, **kwargs):
+        return self.retriever.tokenize(*args, **kwargs)
+
+    def encode(self, *args, **kwargs):
+        return self.retriever.encode(*args, **kwargs)
+
+    def forward(self, *args, **kwargs):
+        return self.retriever(*args, **kwargs)

retriever_config.json

@@ -0,0 +1,135 @@
+{
+  "model_type": "dream",
+  "mask_token_id": 151666,
+  "hidden_size": 3584,
+  "max_length": 156,
+  "n_gen_tokens": 4,
+  "n_gen_q_tokens": 4,
+  "n_gen_p_tokens": 16,
+  "temperature": 0.01,
+  "flops_weight": 0.0,
+  "num_denoise_steps": 1,
+  "sparse_weight": 1.0,
+  "normalize": true,
+  "query_prefix_ids": [
+    151644,
+    8948,
+    198,
+    2610,
+    525,
+    458,
+    15235,
+    17847,
+    429,
+    646,
+    3535,
+    3738,
+    4128,
+    13,
+    151645,
+    198,
+    151644,
+    872,
+    198,
+    2859,
+    25,
+    330
+  ],
+  "query_suffix_ids": [
+    3263,
+    5443,
+    264,
+    2421,
+    4244,
+    311,
+    4009,
+    279,
+    3239,
+    304,
+    264,
+    56370,
+    3383,
+    13,
+    7405,
+    2704,
+    697,
+    4244,
+    525,
+    304,
+    42047,
+    13,
+    151645,
+    198,
+    151644,
+    77091,
+    198,
+    785,
+    4244,
+    525,
+    25,
+    330
+  ],
+  "passage_prefix_ids": [
+    151644,
+    8948,
+    198,
+    2610,
+    525,
+    458,
+    15235,
+    17847,
+    429,
+    646,
+    3535,
+    3738,
+    4128,
+    13,
+    151645,
+    198,
+    151644,
+    872,
+    198,
+    12187,
+    424,
+    25,
+    330
+  ],
+  "passage_suffix_ids": [
+    3263,
+    5443,
+    264,
+    2421,
+    4244,
+    311,
+    4009,
+    279,
+    21085,
+    304,
+    264,
+    56370,
+    3383,
+    13,
+    7405,
+    2704,
+    697,
+    4244,
+    525,
+    304,
+    42047,
+    13,
+    151645,
+    198,
+    151644,
+    77091,
+    198,
+    785,
+    4244,
+    525,
+    25,
+    330
+  ],
+  "lora_rank": 16,
+  "lora_alpha": 64,
+  "use_eos": false,
+  "n_tail": 3
+}

sparse_utils.py

@@ -0,0 +1,93 @@
+"""
+Sparse representation utilities matching original PromptReps.
+
+Reference: https://github.com/ielab/PromptReps
+The original filters sparse logits by:
+1. Word-tokenizing the text (lowercased) with NLTK
+2. Removing stopwords and punctuation
+3. Re-tokenizing each content word independently to get clean token IDs
+4. Only keeping logits for those token IDs
+"""
+
+import torch
+from typing import List, Set, Optional
+import string
+import logging
+
+logger = logging.getLogger(__name__)
+
+try:
+    from nltk import word_tokenize
+    from nltk.corpus import stopwords as _sw_corpus
+    STOPWORDS = set(_sw_corpus.words('english') + list(string.punctuation))
+except LookupError:
+    import nltk
+    nltk.download('punkt_tab', quiet=True)
+    nltk.download('stopwords', quiet=True)
+    from nltk import word_tokenize
+    from nltk.corpus import stopwords as _sw_corpus
+    STOPWORDS = set(_sw_corpus.words('english') + list(string.punctuation))
+
+
+def get_content_token_ids(texts: List[str], tokenizer) -> List[Set[int]]:
+    """Extract content token IDs from texts (stopwords removed, word-level tokenization).
+
+    Matches original PromptReps get_valid_tokens_values:
+    1. word_tokenize(text.lower())
+    2. Remove stopwords + punctuation
+    3. tokenizer.encode(word) for each remaining word
+
+    Args:
+        texts: List of raw text strings.
+        tokenizer: HuggingFace tokenizer.
+
+    Returns:
+        List of sets of token IDs, one per text.
+    """
+    all_token_ids = []
+    for text in texts:
+        words = [w for w in word_tokenize(text.lower())
+                 if w not in STOPWORDS]
+        if words:
+            # Batch encode all words at once (much faster than per-word encode)
+            batch_ids = tokenizer(words, add_special_tokens=False)['input_ids']
+            token_ids = set()
+            for ids in batch_ids:
+                token_ids.update(ids)
+        else:
+            token_ids = set()
+        all_token_ids.append(token_ids)
+    return all_token_ids
+
+
+def filter_sparse(
+    sparse: torch.Tensor,
+    content_token_ids: List[Set[int]],
+    exclude_ids: Optional[List[int]] = None,
+) -> torch.Tensor:
+    """Filter sparse logits to only keep content tokens.
+
+    Args:
+        sparse: [batch, vocab] sparse logit tensor.
+        content_token_ids: List of sets of valid token IDs per example.
+        exclude_ids: Token IDs to always exclude (e.g., MASK token).
+
+    Returns:
+        Filtered sparse tensor (same shape, non-content entries zeroed).
+    """
+    # Build (row, col) index pairs for all content tokens across batch
+    rows, cols = [], []
+    for i in range(sparse.size(0)):
+        if content_token_ids[i]:
+            ids = list(content_token_ids[i])
+            rows.extend([i] * len(ids))
+            cols.extend(ids)
+    if rows:
+        mask = torch.zeros_like(sparse)
+        mask[rows, cols] = 1.0
+    else:
+        mask = torch.zeros_like(sparse)
+    if exclude_ids:
+        for eid in exclude_ids:
+            mask[:, eid] = 0.0
+    return sparse * mask

special_tokens_map.json

@@ -0,0 +1,34 @@
+{
+  "additional_special_tokens": [
+    "<|beginoftext|>",
+    "<|mask|>"
+  ],
+  "bos_token": {
+    "content": "<|beginoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<|mask|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenization_dream.py

@@ -0,0 +1,340 @@
+# coding=utf-8
+# Copyright 2024 The Dream team, HKUNLP Group and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on Qwen's implementations in this library.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for Dream."""
+
+import json
+import os
+import unicodedata
+from functools import lru_cache
+from typing import Optional, Tuple
+
+import regex as re
+
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "merges_file": "merges.txt",
+}
+
+
+MAX_MODEL_INPUT_SIZES = {"dream/dream-tokenizer": 32768}
+
+PRETOKENIZE_REGEX = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+
+
+@lru_cache()
+# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
+    characters the bpe code barfs on.
+
+    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
+    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
+    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
+    tables between utf-8 bytes and unicode strings.
+    """
+    bs = (
+        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs
+def get_pairs(word):
+    """
+    Return set of symbol pairs in a word.
+
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+class DreamTokenizer(PreTrainedTokenizer):
+    """
+    Construct a Dream tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will
+    be encoded differently whether it is at the beginning of the sentence (without space) or not:
+
+    ```python
+    >>> from transformers import AutoTokenizer
+
+    >>> tokenizer = AutoTokenizer.from_pretrained("Dream-org/Dream-v0-Base-7B", trust_remote_code=True)
+    >>> tokenizer("Hello world")["input_ids"]
+    [9707, 1879]
+
+    >>> tokenizer(" Hello world")["input_ids"]
+    [21927, 1879]
+    ```
+    This is expected.
+
+    You should not use GPT2Tokenizer instead, because of the different pretokenization rules.
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        merges_file (`str`):
+            Path to the merges file.
+        errors (`str`, *optional*, defaults to `"replace"`):
+            Paradigm to follow when decoding bytes to UTF-8. See
+            [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not applicable for this tokenizer.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+            Whether or not the model should cleanup the spaces that were added when splitting the input text during the
+            tokenization process. Not applicable to this tokenizer, since tokenization does not add spaces.
+        split_special_tokens (`bool`, *optional*, defaults to `False`):
+            Whether or not the special tokens should be split during the tokenization process. The default behavior is
+            to not split special tokens. This means that if `<|endoftext|>` is the `eos_token`, then `tokenizer.tokenize("<|endoftext|>") =
+            ['<|endoftext|>`]. Otherwise, if `split_special_tokens=True`, then `tokenizer.tokenize("<|endoftext|>")` will be give `['<',
+            '|', 'endo', 'ft', 'ext', '|', '>']`. This argument is only supported for `slow` tokenizers for the moment.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        errors="replace",
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        clean_up_tokenization_spaces=False,
+        split_special_tokens=False,
+        **kwargs,
+    ):
+        # Dream vocab does not contain control tokens; added tokens need to be special
+        bos_token = (
+            AddedToken(bos_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(bos_token, str)
+            else bos_token
+        )
+        eos_token = (
+            AddedToken(eos_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(eos_token, str)
+            else eos_token
+        )
+        unk_token = (
+            AddedToken(unk_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(unk_token, str)
+            else unk_token
+        )
+        pad_token = (
+            AddedToken(pad_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(pad_token, str)
+            else pad_token
+        )
+
+        with open(vocab_file, encoding="utf-8") as vocab_handle:
+            self.encoder = json.load(vocab_handle)
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.errors = errors  # how to handle errors in decoding
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        bpe_merges = []
+        with open(merges_file, encoding="utf-8") as merges_handle:
+            for i, line in enumerate(merges_handle):
+                line = line.strip()
+                if (i == 0 and line.startswith("#version:")) or not line:
+                    continue
+                bpe_merges.append(tuple(line.split()))
+        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        # NOTE: the cache can grow without bound and will get really large for long running processes
+        # (esp. for texts of language that do not use space between word, e.g. Chinese); technically
+        # not a memory leak but appears as one.
+        # GPT2Tokenizer has the same problem, so let's be consistent.
+        self.cache = {}
+
+        self.pat = re.compile(PRETOKENIZE_REGEX)
+
+        if kwargs.get("add_prefix_space", False):
+            logger.warning_once(
+                f"{self.__class__.__name} does not support `add_prefix_space`, setting it to True has no effect."
+            )
+
+        super().__init__(
+            errors=errors,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            split_special_tokens=split_special_tokens,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.encoder)
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab
+    def get_vocab(self):
+        return dict(self.encoder, **self.added_tokens_encoder)
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token
+
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                else:
+                    new_word.extend(word[i:j])
+                    i = j
+
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize
+    def _tokenize(self, text):
+        """Tokenize a string."""
+        bpe_tokens = []
+        for token in re.findall(self.pat, text):
+            token = "".join(
+                self.byte_encoder[b] for b in token.encode("utf-8")
+            )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
+            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
+        return bpe_tokens
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.decoder.get(index)
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        text = "".join(tokens)
+        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
+        return text
+
+    def decode(
+        self,
+        token_ids,
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = False,
+        spaces_between_special_tokens: bool = False,
+        **kwargs,
+    ) -> str:
+        # `spaces_between_special_tokens` defaults to True for _decode in slow tokenizers
+        # and cannot be configured elsewhere, but it should default to False for DreamTokenizer
+        return super().decode(
+            token_ids,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            spaces_between_special_tokens=spaces_between_special_tokens,
+            **kwargs,
+        )
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+        merge_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
+        )
+
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
+
+        index = 0
+        with open(merge_file, "w", encoding="utf-8") as writer:
+            writer.write("#version: 0.2\n")
+            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    logger.warning(
+                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
+                        " Please check that the tokenizer is not corrupted!"
+                    )
+                    index = token_index
+                writer.write(" ".join(bpe_tokens) + "\n")
+                index += 1
+
+        return vocab_file, merge_file
+
+    def prepare_for_tokenization(self, text, **kwargs):
+        text = unicodedata.normalize("NFC", text)
+        return (text, kwargs)

tokenizer_config.json

@@ -0,0 +1,219 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<|beginoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151666": {
+      "content": "<|mask|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|beginoftext|>",
+    "<|mask|>"
+  ],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_dream.DreamTokenizer",
+      null
+    ]
+  },
+  "bos_token": "<|beginoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "mask_token": "<|mask|>",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "DreamTokenizer",
+  "unk_token": null
+}