Add files using upload-large-folder tool

adapters/__init__.py

@@ -0,0 +1 @@
+"""Compatibility imports for the Task 7 adapter paths."""

adapters/_distilled_runtime.py

@@ -0,0 +1,586 @@
+"""Server-only runtime shared by the SDTT and Di4C sampling entrypoints."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+import hashlib
+import json
+import os
+from pathlib import Path
+import random
+import re
+import sys
+import tempfile
+from collections.abc import Mapping
+
+import yaml
+
+from dlb.adapters.conditional_runtime import load_conditioning_batch, token_project_fn
+
+
+class TokenizerBinding:
+    __slots__ = ("tokenizer_id", "revision", "snapshot")
+
+    def __init__(self, tokenizer_id: str, revision: str, snapshot: Path) -> None:
+        self.tokenizer_id = tokenizer_id
+        self.revision = revision
+        self.snapshot = snapshot
+
+
+def require_file(path: Path, label: str) -> Path:
+    path = path.absolute()
+    if path.is_symlink() or not path.is_file() or path.stat().st_size <= 0:
+        raise ValueError(f"{label} is missing or unsafe: {path}")
+    return path
+
+
+def require_directory(path: Path, label: str) -> Path:
+    path = path.absolute()
+    if path.is_symlink() or not path.is_dir():
+        raise ValueError(f"{label} is missing or unsafe: {path}")
+    return path
+
+
+def sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with require_file(path, "hashed file").open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def verify_sha256(path: Path, expected: str, label: str) -> Path:
+    path = require_file(path, label)
+    if (
+        len(expected) != 64
+        or any(character not in "0123456789abcdef" for character in expected)
+        or sha256_file(path) != expected
+    ):
+        raise ValueError(f"{label} SHA-256 differs from canonical provenance")
+    return path
+
+
+def parse_bool(value: str) -> bool:
+    if value == "true":
+        return True
+    if value == "false":
+        return False
+    raise ValueError("boolean values must be 'true' or 'false'")
+
+
+@contextmanager
+def offline_huggingface(enabled: bool):
+    names = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")
+    previous = {name: os.environ.get(name) for name in names}
+    if enabled:
+        os.environ.update({name: "1" for name in names})
+    try:
+        yield
+    finally:
+        for name, value in previous.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value
+
+
+def install_upstream(upstream_root: Path) -> None:
+    source = require_directory(upstream_root / "src", "upstream Python source")
+    source_text = str(source)
+    if source_text not in sys.path:
+        sys.path.insert(0, source_text)
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    try:
+        import numpy
+
+        numpy.random.seed(seed % (2**32))
+    except ImportError:
+        pass
+    import torch
+
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def load_tokenizer(snapshot: Path):
+    from transformers import AutoTokenizer
+
+    return AutoTokenizer.from_pretrained(
+        str(require_directory(snapshot, "locked tokenizer snapshot")),
+        local_files_only=True,
+        trust_remote_code=False,
+    )
+
+
+def load_config(path: Path, expected_sha256: str):
+    from omegaconf import OmegaConf
+
+    return OmegaConf.load(verify_sha256(path, expected_sha256, "sampling config"))
+
+
+def checkpoint_state(
+    path: Path, expected_sha256: str
+) -> tuple[dict[str, object], object | None]:
+    path = verify_sha256(path, expected_sha256, "student checkpoint")
+    if path.suffix == ".safetensors":
+        from safetensors.torch import load_file
+
+        return dict(load_file(str(path))), None
+    import torch
+
+    # Lightning 2.2 checkpoints contain OmegaConf DictConfig objects. Full pickle
+    # loading is deliberately restricted to bytes authenticated above by the
+    # canonical runner/checkpoint manifest provenance.
+    payload = torch.load(str(path), map_location="cpu", weights_only=False)
+    if not isinstance(payload, dict):
+        raise ValueError("student checkpoint is not a mapping")
+    state = payload.get("state_dict", payload)
+    if not isinstance(state, dict):
+        raise ValueError("student checkpoint state_dict is invalid")
+    config = None
+    hyperparameters = payload.get("hyper_parameters")
+    if isinstance(hyperparameters, dict):
+        config = hyperparameters.get("config")
+    return dict(state), config
+
+
+_MISSING = object()
+_STANDARD_REFERENCE = re.compile(r"^\$\{([A-Za-z_][A-Za-z0-9_.]*)\}$")
+
+
+def _select_config_value(
+    config: object, path: str, references: frozenset[str] = frozenset()
+) -> object:
+    """Resolve one selected scalar without walking unrelated Hydra config nodes."""
+
+    try:
+        from omegaconf import OmegaConf
+
+        if OmegaConf.is_config(config):
+            try:
+                return OmegaConf.select(
+                    config,
+                    path,
+                    default=_MISSING,
+                    throw_on_resolution_failure=True,
+                )
+            except Exception as error:
+                raise ValueError(
+                    f"sampling config field {path} cannot be resolved"
+                ) from error
+    except ImportError:
+        pass
+    value = config
+    for component in path.split("."):
+        try:
+            value = (
+                value[component]
+                if isinstance(value, Mapping)
+                else getattr(value, component)
+            )
+        except (AttributeError, KeyError, TypeError):
+            return _MISSING
+    if isinstance(value, str):
+        reference = _STANDARD_REFERENCE.fullmatch(value)
+        if reference is not None:
+            target = reference.group(1)
+            if target in references or target == path:
+                raise ValueError(f"sampling config has a cyclic reference at {path}")
+            return _select_config_value(config, target, references | {path})
+    return value
+
+
+def _nested(config: object, path: str) -> object:
+    value = _select_config_value(config, path)
+    if value is _MISSING:
+        raise ValueError(f"sampling config is missing required field {path}")
+    return value
+
+
+def _optional_nested(config: object, path: str) -> object:
+    return _select_config_value(config, path)
+
+
+_CHECKPOINT_CONFIG_FIELDS = (
+    "T",
+    "time_conditioning",
+    "model.type",
+    "model.hidden_size",
+    "model.cond_dim",
+    "model.length",
+    "model.n_blocks",
+    "model.n_heads",
+    "model.scale_by_sigma",
+    "model.dropout",
+    "model.tie_word_embeddings",
+    "noise.type",
+    "noise.sigma_min",
+    "noise.sigma_max",
+    "training.ema",
+    "training.antithetic_sampling",
+    "training.importance_sampling",
+    "training.sampling_eps",
+    "training.change_of_variables",
+    "tokenizer.name",
+    "parameterization.name",
+    "parameterization.log_loss_buckets",
+    "parameterization.start_from_hf",
+    "parameterization.distill_mode",
+    "parameterization.num_distill_steps",
+    "parameterization.min_num_sampling_steps",
+    "parameterization.grow_dt_every",
+    "parameterization.orig_num_sampling_steps",
+    "parameterization.sampling_mode",
+    "parameterization.loss_precision",
+    "parameterization.reset_optimizer_on_growth",
+    "parameterization.use_ema_on_growth",
+)
+
+_OPTIONAL_CHECKPOINT_CONFIG_FIELDS = ("model.causal", "is_di4c")
+
+
+def validate_embedded_config(
+    authoritative: object,
+    embedded: object | None,
+    *,
+    allow_missing_embedded_config: bool = False,
+) -> None:
+    if embedded is None:
+        if allow_missing_embedded_config:
+            return
+        raise ValueError("Lightning checkpoint does not embed its training config")
+    for path in (*_CHECKPOINT_CONFIG_FIELDS, *_OPTIONAL_CHECKPOINT_CONFIG_FIELDS):
+        expected = _optional_nested(authoritative, path)
+        if expected is _MISSING:
+            continue
+        observed = _optional_nested(embedded, path)
+        if observed is _MISSING:
+            raise ValueError(f"embedded checkpoint config is missing required field {path}")
+        if observed != expected:
+            raise ValueError(
+                f"embedded checkpoint config differs at {path}: "
+                f"expected {expected!r}, found {observed!r}"
+            )
+
+
+def validate_sampling_config(
+    config: object,
+    *,
+    binding: TokenizerBinding,
+    sequence_length: int,
+    require_di4c: bool,
+) -> None:
+    for path in _CHECKPOINT_CONFIG_FIELDS:
+        _nested(config, path)
+    tokenizer_name = _nested(config, "tokenizer.name")
+    if tokenizer_name != binding.tokenizer_id and Path(str(tokenizer_name)).resolve() != binding.snapshot.resolve():
+        raise ValueError("sampling config tokenizer differs from canonical data contract")
+    if _nested(config, "model.length") != sequence_length:
+        raise ValueError("sampling config sequence length differs from canonical dataset")
+    if _nested(config, "data_preprocess.seq_len") != sequence_length:
+        raise ValueError("sampling config preprocessing length differs from canonical dataset")
+    if _nested(config, "parameterization.sampling_mode") != "ancestral":
+        raise ValueError("sampling config does not use ancestral sampling")
+    di4c_identity = _optional_nested(config, "is_di4c")
+    if require_di4c:
+        if di4c_identity is not True:
+            raise ValueError("sampling config Di4C identity is inconsistent")
+        if _nested(config, "parameterization.start_from_hf") is not False:
+            raise ValueError("Di4C sampling config must instantiate the local student")
+    elif di4c_identity not in {_MISSING, False}:
+        raise ValueError("sampling config Di4C identity is inconsistent")
+
+
+def load_tokenizer_binding(
+    data_config_path: Path,
+    downloads_manifest_path: Path,
+    dataset_id: str,
+    requested_snapshot: Path,
+) -> TokenizerBinding:
+    data_config_path = require_file(data_config_path, "data config")
+    downloads_manifest_path = require_file(
+        downloads_manifest_path, "download manifest"
+    )
+    try:
+        configuration = yaml.safe_load(data_config_path.read_text(encoding="utf-8"))
+        downloads = json.loads(downloads_manifest_path.read_text(encoding="utf-8"))
+        tokenizer_id = configuration["datasets"][dataset_id]["tokenizer"]
+        revision = configuration["models"][tokenizer_id]
+        record = downloads["models"][tokenizer_id]
+    except (OSError, UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError) as error:
+        raise ValueError("locked tokenizer metadata is invalid") from error
+    if downloads.get("schema_version") != 1:
+        raise ValueError("download manifest schema is not supported")
+    if (
+        not isinstance(tokenizer_id, str)
+        or not isinstance(revision, str)
+        or len(revision) != 40
+    ):
+        raise ValueError("data config tokenizer binding is invalid")
+    if record.get("repo_id") != tokenizer_id or record.get("revision") != revision:
+        raise ValueError("download manifest tokenizer binding differs from data config")
+    recorded = record.get("snapshot_path")
+    if not isinstance(recorded, str) or not recorded:
+        raise ValueError("download manifest tokenizer snapshot is missing")
+    snapshot = Path(recorded)
+    if not snapshot.is_absolute():
+        snapshot = data_config_path.parents[1] / snapshot
+    requested_snapshot = requested_snapshot.absolute()
+    if requested_snapshot != snapshot:
+        raise ValueError("requested tokenizer snapshot differs from download manifest")
+    parts = snapshot.parts
+    try:
+        recorded_revision = parts[parts.index("snapshots") + 1]
+    except (ValueError, IndexError) as error:
+        raise ValueError("download manifest tokenizer snapshot path is invalid") from error
+    if recorded_revision != revision:
+        raise ValueError("download manifest tokenizer snapshot revision differs from data config")
+    require_directory(snapshot, "locked tokenizer snapshot")
+    return TokenizerBinding(tokenizer_id, revision, snapshot)
+
+
+def instantiate_without_training_teacher(model_type, config, tokenizer):
+    """Instantiate the released student architecture without loading a network teacher."""
+
+    original = model_type.prepare_teacher_and_student
+
+    def sampling_only(self):
+        self.teacher = []
+
+    model_type.prepare_teacher_and_student = sampling_only
+    try:
+        return model_type(config, tokenizer, verbose=False)
+    finally:
+        model_type.prepare_teacher_and_student = original
+
+
+def normalize_state_dict(state: dict[str, object]) -> dict[str, object]:
+    return {key.replace("_orig_mod.", ""): value for key, value in state.items()}
+
+
+def write_capture_atomic(
+    output: Path,
+    *,
+    model,
+    tokenizer,
+    sample_count: int,
+    batch_size: int,
+    num_steps: int,
+    seq_len: int,
+    sampler: str,
+    generation_mode: str = "unconditional",
+    conditioning_manifest: Path | None = None,
+    conditioning_manifest_sha256: str | None = None,
+    prompt_count: int | None = None,
+    diversity_prompt_count: int | None = None,
+    completions_per_diversity_prompt: int | None = None,
+) -> None:
+    """Sample exact batches, immediately move IDs to CPU, and atomically publish JSON."""
+
+    if min(sample_count, batch_size, num_steps, seq_len) <= 0:
+        raise ValueError("sampling counts and lengths must be positive")
+    if sampler != "ancestral":
+        raise ValueError("distilled language baselines require ancestral sampling")
+    if generation_mode not in {"unconditional", "conditional_prefix"}:
+        raise ValueError(f"unsupported generation mode: {generation_mode!r}")
+    output = output.absolute()
+    if not output.is_absolute() or output.is_symlink():
+        raise ValueError(f"output path is unsafe: {output}")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, temporary_name = tempfile.mkstemp(
+        prefix=f".{output.name}.", suffix=".partial", dir=output.parent
+    )
+    temporary = Path(temporary_name)
+    written = 0
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+            handle.write('{"schema":"dlb-upstream-token-capture-v1","samples":[')
+            if generation_mode == "conditional_prefix":
+                if (
+                    conditioning_manifest is None
+                    or conditioning_manifest_sha256 is None
+                    or prompt_count is None
+                    or diversity_prompt_count is None
+                    or completions_per_diversity_prompt is None
+                ):
+                    raise ValueError("conditional distilled sampling requires a prompt contract")
+                batches: list[tuple[int, int, int]] = []
+                for start in range(0, prompt_count, batch_size):
+                    batches.append((0, start, min(batch_size, prompt_count - start)))
+                for completion_id in range(1, completions_per_diversity_prompt):
+                    for start in range(0, diversity_prompt_count, batch_size):
+                        batches.append(
+                            (
+                                completion_id,
+                                start,
+                                min(batch_size, diversity_prompt_count - start),
+                            )
+                        )
+            else:
+                batches = [(-1, written, min(batch_size, sample_count - written)) for written in range(0, sample_count, batch_size)]
+            for completion_id, prompt_start, current in batches:
+                current = min(batch_size, sample_count - written)
+                batch = None
+                project_fn = lambda value: value
+                if generation_mode == "conditional_prefix":
+                    batch = load_conditioning_batch(
+                        conditioning_manifest,
+                        conditioning_manifest_sha256,
+                        completion_id=completion_id,
+                        prompt_start=prompt_start,
+                        batch_size=current,
+                        device=str(model.device),
+                        vocab_size=len(tokenizer),
+                    )
+                    project_fn = token_project_fn(batch.prefix_token_ids)
+                result = project_fn(
+                    model.sample(
+                        n_samples=current,
+                        num_steps=num_steps,
+                        seq_len=seq_len,
+                        sampler=sampler,
+                        verbose=False,
+                        project_fn=project_fn,
+                    )
+                )
+                rows = result.detach().cpu().tolist()
+                if not isinstance(rows, list) or len(rows) != current:
+                    raise ValueError("upstream sampler returned the wrong batch size")
+                texts = list(tokenizer.batch_decode(rows))
+                if len(texts) != current:
+                    raise ValueError("tokenizer returned the wrong decoded batch size")
+                for batch_index, (row, text) in enumerate(zip(rows, texts, strict=True)):
+                    if not isinstance(row, list) or len(row) != seq_len:
+                        raise ValueError("upstream sampler returned a noncanonical sequence")
+                    if not isinstance(text, str) or not text.strip():
+                        raise ValueError("upstream tokenizer returned empty text")
+                    payload = {"sample_id": written, "text": text, "token_ids": row}
+                    if batch is not None:
+                        prefix = batch.prefix_token_ids[batch_index].detach().cpu().tolist()
+                        reference = batch.reference_token_ids[batch_index].detach().cpu().tolist()
+                        if row[: len(prefix)] != prefix:
+                            raise ValueError("conditional sampler returned a prefix mismatch")
+                        payload.update(
+                            {
+                                "prompt_id": batch.prompt_ids[batch_index],
+                                "completion_id": completion_id,
+                                "source_index": batch.source_indices[batch_index],
+                                "prefix_token_ids": prefix,
+                                "reference_token_ids": reference,
+                                "full_token_ids": row,
+                                "prefix_text": tokenizer.decode(prefix),
+                                "continuation_text": tokenizer.decode(row[len(prefix) :]),
+                                "reference_text": tokenizer.decode(reference),
+                                "full_text": text,
+                                "prefix_exact_match": True,
+                            }
+                        )
+                    if written:
+                        handle.write(",")
+                    json.dump(
+                        payload,
+                        handle,
+                        ensure_ascii=False,
+                        separators=(",", ":"),
+                    )
+                    written += 1
+                del result, rows, texts
+            handle.write("]}")
+            handle.flush()
+            os.fsync(handle.fileno())
+        if written != sample_count:
+            raise ValueError(f"expected {sample_count} samples, wrote {written}")
+        os.replace(temporary, output)
+    except BaseException:
+        temporary.unlink(missing_ok=True)
+        raise
+
+
+def benchmark_model(
+    *,
+    model,
+    output: Path,
+    metadata_path: Path,
+    precision: str,
+    num_steps: int,
+    seq_len: int,
+    sampler: str,
+    generation_mode: str = "unconditional",
+    conditioning_manifest: Path | None = None,
+    conditioning_manifest_sha256: str | None = None,
+    prompt_count: int | None = None,
+    diversity_prompt_count: int | None = None,
+    completions_per_diversity_prompt: int | None = None,
+):
+    """Time the already-materialized student's real one-sample tensor callback."""
+
+    if min(num_steps, seq_len) <= 0 or sampler != "ancestral":
+        raise ValueError("distilled benchmark sampling arguments are invalid")
+    from dlb.timing import benchmark_and_publish
+
+    project_fn = lambda value: value
+    if generation_mode == "conditional_prefix":
+        if conditioning_manifest is None or conditioning_manifest_sha256 is None:
+            raise ValueError("conditional benchmark requires a prompt manifest")
+        del prompt_count, diversity_prompt_count, completions_per_diversity_prompt
+        batch = load_conditioning_batch(
+            conditioning_manifest,
+            conditioning_manifest_sha256,
+            completion_id=0,
+            prompt_start=0,
+            batch_size=1,
+            device=str(model.device),
+            vocab_size=len(model.tokenizer),
+        )
+        project_fn = token_project_fn(batch.prefix_token_ids)
+    return benchmark_and_publish(
+        lambda: project_fn(model.sample(
+            n_samples=1,
+            num_steps=num_steps,
+            seq_len=seq_len,
+            sampler=sampler,
+            verbose=False,
+            project_fn=project_fn,
+        )),
+        model=model,
+        output=output,
+        metadata_path=metadata_path,
+        precision=precision,
+    )
+
+
+def configure_for_sampling(config, *, tokenizer_snapshot: Path, seq_len: int) -> None:
+    config.tokenizer.name = str(tokenizer_snapshot)
+    config.model.length = seq_len
+    config.data_preprocess.seq_len = seq_len
+    config.mode = "sample"
+    config.compile = False
+
+
+def materialize_model(
+    *,
+    model_type,
+    config,
+    tokenizer,
+    state: dict[str, object],
+    strict: bool,
+):
+    model = instantiate_without_training_teacher(model_type, config, tokenizer)
+    result = model.load_state_dict(normalize_state_dict(state), strict=strict)
+    if not strict:
+        missing = [key for key in result.missing_keys if not key.startswith("teacher.")]
+        unexpected = [key for key in result.unexpected_keys if not key.startswith("teacher.")]
+        if missing or unexpected:
+            raise ValueError(
+                "student checkpoint does not match the configured architecture: "
+                f"missing={missing[:8]}, unexpected={unexpected[:8]}"
+            )
+    model.eval()
+    if not __import__("torch").cuda.is_available():
+        raise RuntimeError("distilled sampling requires a CUDA server")
+    model.cuda()
+    return model

adapters/base.py

@@ -0,0 +1,3 @@
+from dlb.adapters.base import AdapterError, BaseTeacherAdapter, CheckpointSelection
+
+__all__ = ["AdapterError", "BaseTeacherAdapter", "CheckpointSelection"]

adapters/candi.py

@@ -0,0 +1,3 @@
+from dlb.adapters.candi import CANDIAdapter, adapter
+
+__all__ = ["CANDIAdapter", "adapter"]

adapters/di4c.py

@@ -0,0 +1,3 @@
+from dlb.adapters.di4c import Di4CAdapter, adapter
+
+__all__ = ["Di4CAdapter", "adapter"]

adapters/duo.py

@@ -0,0 +1,3 @@
+from dlb.adapters.duo import DuoAdapter, adapter
+
+__all__ = ["DuoAdapter", "adapter"]

adapters/flm.py

@@ -0,0 +1,3 @@
+from dlb.adapters.flm import FLMAdapter, adapter
+
+__all__ = ["FLMAdapter", "adapter"]

adapters/langflow.py

@@ -0,0 +1,3 @@
+from dlb.adapters.langflow import LangFlowAdapter, adapter
+
+__all__ = ["LangFlowAdapter", "adapter"]

adapters/mdlm.py

@@ -0,0 +1,3 @@
+from dlb.adapters.mdlm import MDLMAdapter, adapter
+
+__all__ = ["MDLMAdapter", "adapter"]

adapters/rdlm.py

@@ -0,0 +1,3 @@
+from dlb.adapters.rdlm import RDLMAdapter, adapter
+
+__all__ = ["RDLMAdapter", "adapter"]

adapters/sample_di4c.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Sample a pinned Di4C language student from local server assets only."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from _distilled_runtime import (
+    benchmark_model,
+    checkpoint_state,
+    configure_for_sampling,
+    install_upstream,
+    load_config,
+    load_tokenizer,
+    load_tokenizer_binding,
+    materialize_model,
+    offline_huggingface,
+    parse_bool,
+    require_directory,
+    require_file,
+    seed_everything,
+    validate_embedded_config,
+    validate_sampling_config,
+    write_capture_atomic,
+)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--upstream-root", type=Path, required=True)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument("--checkpoint-sha256", required=True)
+    parser.add_argument("--config-sha256", required=True)
+    parser.add_argument("--data-config", type=Path, required=True)
+    parser.add_argument("--downloads-manifest", type=Path, required=True)
+    parser.add_argument("--tokenizer-snapshot", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--sample-count", type=int, required=True)
+    parser.add_argument("--batch-size", type=int, required=True)
+    parser.add_argument("--num-steps", type=int, required=True)
+    parser.add_argument("--seq-len", type=int, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--sampler", choices=("ancestral",), required=True)
+    parser.add_argument(
+        "--teacher-family", choices=("masked_mdlm", "uniform_duo"), required=True
+    )
+    parser.add_argument("--dataset", choices=("lm1b", "owt"), required=True)
+    parser.add_argument("--generation-mode", choices=("unconditional", "conditional_prefix"), default="unconditional")
+    parser.add_argument("--conditioning-manifest", type=Path)
+    parser.add_argument("--conditioning-manifest-sha256")
+    parser.add_argument("--conditioning-config-sha256")
+    parser.add_argument("--prefix-length", type=int)
+    parser.add_argument("--evaluation-continuation-length", type=int)
+    parser.add_argument("--prompt-count", type=int)
+    parser.add_argument("--diversity-prompt-count", type=int)
+    parser.add_argument("--completions-per-diversity-prompt", type=int)
+    parser.add_argument("--completion-schedule")
+    parser.add_argument("--offline", type=parse_bool, required=True)
+    parser.add_argument(
+        "--allow-missing-embedded-config", type=parse_bool, default=False
+    )
+    parser.add_argument("--benchmark-output", type=Path)
+    parser.add_argument("--benchmark-metadata", type=Path)
+    parser.add_argument("--benchmark-precision", choices=("author",))
+    args = parser.parse_args(argv)
+    benchmark_values = (args.benchmark_output, args.benchmark_metadata, args.benchmark_precision)
+    if any(value is not None for value in benchmark_values) and not all(
+        value is not None for value in benchmark_values
+    ):
+        parser.error("benchmark arguments must be provided together")
+
+    upstream = require_directory(args.upstream_root, "pinned Di4C language source")
+    require_file(upstream / "src/sdtt/main.py", "pinned Di4C entrypoint")
+    install_upstream(upstream)
+    checkpoint = require_file(args.checkpoint, "Di4C student checkpoint")
+    config_path = require_file(args.config, "Di4C sampling config")
+    binding = load_tokenizer_binding(
+        args.data_config,
+        args.downloads_manifest,
+        args.dataset,
+        args.tokenizer_snapshot,
+    )
+    with offline_huggingface(args.offline):
+        from sdtt.core.distill.multi_round_sdtt import MultiRoundSDTT
+
+        config = load_config(config_path, args.config_sha256)
+        state, embedded_config = checkpoint_state(
+            checkpoint, args.checkpoint_sha256
+        )
+        validate_sampling_config(
+            config,
+            binding=binding,
+            sequence_length=args.seq_len,
+            require_di4c=True,
+        )
+        validate_embedded_config(
+            config,
+            embedded_config,
+            allow_missing_embedded_config=args.allow_missing_embedded_config,
+        )
+        configure_for_sampling(
+            config, tokenizer_snapshot=binding.snapshot, seq_len=args.seq_len
+        )
+        config.parameterization.checkpoint_path = str(checkpoint)
+        tokenizer = load_tokenizer(binding.snapshot)
+        seed_everything(args.seed)
+        model = materialize_model(
+            model_type=MultiRoundSDTT,
+            config=config,
+            tokenizer=tokenizer,
+            state=state,
+            strict=False,
+        )
+        if args.benchmark_output is not None:
+            benchmark_model(
+                model=model,
+                output=args.benchmark_output,
+                metadata_path=args.benchmark_metadata,
+                precision=args.benchmark_precision,
+                num_steps=args.num_steps,
+                seq_len=args.seq_len,
+                sampler=args.sampler,
+                generation_mode=args.generation_mode,
+                conditioning_manifest=args.conditioning_manifest,
+                conditioning_manifest_sha256=args.conditioning_manifest_sha256,
+                prompt_count=args.prompt_count,
+                diversity_prompt_count=args.diversity_prompt_count,
+                completions_per_diversity_prompt=args.completions_per_diversity_prompt,
+            )
+        else:
+            write_capture_atomic(
+                args.output,
+                model=model,
+                tokenizer=tokenizer,
+                sample_count=args.sample_count,
+                batch_size=args.batch_size,
+                num_steps=args.num_steps,
+                seq_len=args.seq_len,
+                sampler=args.sampler,
+                generation_mode=args.generation_mode,
+                conditioning_manifest=args.conditioning_manifest,
+                conditioning_manifest_sha256=args.conditioning_manifest_sha256,
+                prompt_count=args.prompt_count,
+                diversity_prompt_count=args.diversity_prompt_count,
+                completions_per_diversity_prompt=args.completions_per_diversity_prompt,
+            )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

adapters/sample_langflow.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+"""Dataset-configurable LangFlow sampler used by the DLB capture wrapper."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoTokenizer
+
+
+ROOT = Path(__file__).resolve().parents[1]
+UPSTREAM_ROOT = ROOT / "upstreams" / "langflow"
+if str(UPSTREAM_ROOT) not in sys.path:
+    sys.path.insert(0, str(UPSTREAM_ROOT))
+
+from langflow import LangFlow, LangFlowConfig  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate samples with pinned LangFlow")
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--num_samples", type=int, default=5)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--num_steps", type=int, default=128)
+    parser.add_argument("--seq_length", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--output")
+    parser.add_argument("--tokenizer", required=True)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    checkpoint = Path(args.checkpoint)
+    config_dir = checkpoint.parent
+
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    config = LangFlowConfig.from_pretrained(str(config_dir))
+    model = LangFlow(config)
+    state_dict = load_file(str(checkpoint), device=str(device))
+    model.load_state_dict(state_dict)
+    model = model.to(device)
+    model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+
+    all_samples = []
+    remaining = args.num_samples
+    with torch.no_grad():
+        while remaining > 0:
+            batch = min(args.batch_size, remaining)
+            samples = model.generate_samples(
+                num_samples=batch,
+                seq_length=args.seq_length,
+                num_steps=args.num_steps,
+                device=device,
+            )
+            all_samples.append(samples)
+            remaining -= batch
+
+    generated = torch.cat(all_samples, dim=0)
+    texts = tokenizer.batch_decode(generated, skip_special_tokens=True)
+    for index, text in enumerate(texts):
+        print(f"\n--- Sample {index + 1} ---")
+        print(text[:500] + ("..." if len(text) > 500 else ""))
+
+    if args.output is not None:
+        with open(args.output, "w", encoding="utf-8") as output_file:
+            for index, text in enumerate(texts):
+                output_file.write(f"--- Sample {index + 1} ---\n{text}\n\n")
+
+
+if __name__ == "__main__":
+    main()

adapters/sample_sdtt.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""Sample the pinned SDTT KLD round-7 student from local server assets only."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from _distilled_runtime import (
+    benchmark_model,
+    configure_for_sampling,
+    install_upstream,
+    load_config,
+    load_tokenizer,
+    load_tokenizer_binding,
+    materialize_model,
+    offline_huggingface,
+    parse_bool,
+    require_directory,
+    require_file,
+    seed_everything,
+    checkpoint_state,
+    validate_embedded_config,
+    validate_sampling_config,
+    write_capture_atomic,
+)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--upstream-root", type=Path, required=True)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument("--checkpoint-sha256", required=True)
+    parser.add_argument("--config-sha256", required=True)
+    parser.add_argument("--data-config", type=Path, required=True)
+    parser.add_argument("--downloads-manifest", type=Path, required=True)
+    parser.add_argument("--dataset", choices=("lm1b", "owt"), required=True)
+    parser.add_argument("--tokenizer-snapshot", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--sample-count", type=int, required=True)
+    parser.add_argument("--batch-size", type=int, required=True)
+    parser.add_argument("--num-steps", type=int, required=True)
+    parser.add_argument("--seq-len", type=int, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--sampler", choices=("ancestral",), required=True)
+    parser.add_argument("--loss", choices=("kld",), required=True)
+    parser.add_argument("--round", type=int, choices=(7,), required=True)
+    parser.add_argument("--teacher-family", choices=("masked_mdlm",), required=True)
+    parser.add_argument("--generation-mode", choices=("unconditional", "conditional_prefix"), default="unconditional")
+    parser.add_argument("--conditioning-manifest", type=Path)
+    parser.add_argument("--conditioning-manifest-sha256")
+    parser.add_argument("--conditioning-config-sha256")
+    parser.add_argument("--prefix-length", type=int)
+    parser.add_argument("--evaluation-continuation-length", type=int)
+    parser.add_argument("--prompt-count", type=int)
+    parser.add_argument("--diversity-prompt-count", type=int)
+    parser.add_argument("--completions-per-diversity-prompt", type=int)
+    parser.add_argument("--completion-schedule")
+    parser.add_argument("--offline", type=parse_bool, required=True)
+    parser.add_argument("--benchmark-output", type=Path)
+    parser.add_argument("--benchmark-metadata", type=Path)
+    parser.add_argument("--benchmark-precision", choices=("author",))
+    args = parser.parse_args(argv)
+    benchmark_values = (args.benchmark_output, args.benchmark_metadata, args.benchmark_precision)
+    if any(value is not None for value in benchmark_values) and not all(
+        value is not None for value in benchmark_values
+    ):
+        parser.error("benchmark arguments must be provided together")
+
+    upstream = require_directory(args.upstream_root, "pinned SDTT source")
+    require_file(upstream / "src/sdtt/main.py", "pinned SDTT entrypoint")
+    install_upstream(upstream)
+    checkpoint = require_file(args.checkpoint, "SDTT student checkpoint")
+    config_path = require_file(args.config, "SDTT student config")
+    binding = load_tokenizer_binding(
+        args.data_config,
+        args.downloads_manifest,
+        args.dataset,
+        args.tokenizer_snapshot,
+    )
+    with offline_huggingface(args.offline):
+        from sdtt.core.distill.multi_round_sdtt import MultiRoundSDTT
+
+        config = load_config(config_path, args.config_sha256)
+        state, embedded_config = checkpoint_state(
+            checkpoint, args.checkpoint_sha256
+        )
+        validate_sampling_config(
+            config,
+            binding=binding,
+            sequence_length=args.seq_len,
+            require_di4c=False,
+        )
+        if embedded_config is not None:
+            validate_embedded_config(config, embedded_config)
+        configure_for_sampling(
+            config, tokenizer_snapshot=binding.snapshot, seq_len=args.seq_len
+        )
+        tokenizer = load_tokenizer(binding.snapshot)
+        seed_everything(args.seed)
+        model = materialize_model(
+            model_type=MultiRoundSDTT,
+            config=config,
+            tokenizer=tokenizer,
+            state=state,
+            strict=True,
+        )
+        if args.benchmark_output is not None:
+            benchmark_model(
+                model=model,
+                output=args.benchmark_output,
+                metadata_path=args.benchmark_metadata,
+                precision=args.benchmark_precision,
+                num_steps=args.num_steps,
+                seq_len=args.seq_len,
+                sampler=args.sampler,
+                generation_mode=args.generation_mode,
+                conditioning_manifest=args.conditioning_manifest,
+                conditioning_manifest_sha256=args.conditioning_manifest_sha256,
+                prompt_count=args.prompt_count,
+                diversity_prompt_count=args.diversity_prompt_count,
+                completions_per_diversity_prompt=args.completions_per_diversity_prompt,
+            )
+        else:
+            write_capture_atomic(
+                args.output,
+                model=model,
+                tokenizer=tokenizer,
+                sample_count=args.sample_count,
+                batch_size=args.batch_size,
+                num_steps=args.num_steps,
+                seq_len=args.seq_len,
+                sampler=args.sampler,
+                generation_mode=args.generation_mode,
+                conditioning_manifest=args.conditioning_manifest,
+                conditioning_manifest_sha256=args.conditioning_manifest_sha256,
+                prompt_count=args.prompt_count,
+                diversity_prompt_count=args.diversity_prompt_count,
+                completions_per_diversity_prompt=args.completions_per_diversity_prompt,
+            )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

adapters/sdtt.py

@@ -0,0 +1,3 @@
+from dlb.adapters.sdtt import SDTTAdapter, adapter
+
+__all__ = ["SDTTAdapter", "adapter"]

adapters/train_di4c.py

@@ -0,0 +1,179 @@
+"""Run Di4C language distillation with project-controlled student initialization."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import sys
+from typing import Any
+
+
+_INIT_MODES = frozenset({"hf_small", "teacher", "scratch"})
+_EVAL_GLOBALS = {
+    "__builtins__": {},
+    "__import__": __import__,
+    "abs": abs,
+    "float": float,
+    "int": int,
+    "len": len,
+    "max": max,
+    "min": min,
+    "round": round,
+    "str": str,
+    "sum": sum,
+}
+
+
+def _install_upstream_path() -> Path:
+    source = Path.cwd()
+    sdtt_src = source / "sdtt" / "src"
+    sdtt_package_root = sdtt_src / "sdtt"
+    if not (sdtt_package_root / "main.py").is_file():
+        raise RuntimeError(f"Di4C SDTT source is missing: {sdtt_src}")
+    for path in (source, sdtt_src, sdtt_package_root):
+        text = str(path)
+        if text not in sys.path:
+            sys.path.insert(0, text)
+    return source
+
+
+def _student_init_mode(config: Any) -> str:
+    mode = config.get("dlb_student_init", "hf_small")
+    if mode not in _INIT_MODES:
+        raise ValueError(
+            f"unsupported dlb_student_init={mode!r}; expected one of {sorted(_INIT_MODES)}"
+        )
+    return str(mode)
+
+
+def _eval_expression(expression: Any) -> Any:
+    return eval(str(expression), _EVAL_GLOBALS, {})
+
+
+def _device_count() -> int:
+    try:
+        import torch
+    except ImportError:
+        return 0
+    return int(torch.cuda.device_count())
+
+
+def _register_omegaconf_resolvers() -> None:
+    from omegaconf import OmegaConf
+
+    resolvers = {
+        "cwd": lambda: str(Path.cwd()),
+        "device_count": _device_count,
+        "div_up": lambda numerator, denominator: (
+            int(numerator) + int(denominator) - 1
+        )
+        // int(denominator),
+        "eval": _eval_expression,
+    }
+    for name, resolver in resolvers.items():
+        try:
+            OmegaConf.register_new_resolver(name, resolver, replace=True)
+        except TypeError:
+            OmegaConf.clear_resolver(name)
+            OmegaConf.register_new_resolver(name, resolver)
+
+
+def _teacher_initialized_student(config: Any):
+    from sdtt.loading_utils import get_diffusion
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(config.tokenizer.name)
+    return get_diffusion(config, tokenizer)
+
+
+def _hf_small_initialized_student(loss: str, round: int, config: Any):
+    if round not in range(1, 8):
+        raise ValueError(
+            f"Round value is too large: should be 1 <= round <= 7. Actual value: `{round}`"
+        )
+    if loss not in ("kld", "mse", "tvd"):
+        raise ValueError(f"Valid losses sizes: kld, mse, tvd. Actual value: `{loss}`")
+
+    from sdtt.core.distill.multi_round_sdtt import MultiRoundSDTT
+
+    revision = f"baselines_{loss}_step_{round * 10_000}"
+    # Upstream load_small_student uses student_as_teacher=True, which overwrites
+    # the configured teacher checkpoint with the HF small student backbone.
+    return MultiRoundSDTT.from_pretrained(
+        "jdeschena/sdtt",
+        revision,
+        config=config,
+        student_as_teacher=False,
+    )
+
+
+def _scratch_initialized_student(config: Any):
+    from sdtt.models.loading_utils import get_backbone
+
+    model = _teacher_initialized_student(config)
+    model.backbone = get_backbone(config, model.vocab_size)
+    if hasattr(model.backbone, "is_di4c"):
+        model.backbone.is_di4c = True
+    model.init_ema()
+    return model
+
+
+def _patch_student_loader() -> None:
+    import sdtt
+
+    original = sdtt.load_small_student
+    if getattr(original, "_dlb_student_init_patch", False):
+        return
+
+    def load_small_student(
+        loss: str = "kld",
+        round: int = 1,
+        config: Any | None = None,
+        **kwargs: Any,
+    ):
+        if config is None:
+            try:
+                return original(loss=loss, round=round, config=config, **kwargs)
+            except TypeError:
+                if kwargs:
+                    return original(loss=loss, round=round, config=config)
+                raise
+
+        mode = _student_init_mode(config)
+        if mode == "hf_small":
+            kwargs.pop("student_as_teacher", None)
+            if kwargs:
+                raise TypeError(
+                    f"unexpected load_small_student keyword(s): {sorted(kwargs)}"
+                )
+            return _hf_small_initialized_student(loss=loss, round=round, config=config)
+        if mode == "teacher":
+            return _teacher_initialized_student(config)
+        if mode == "scratch":
+            return _scratch_initialized_student(config)
+        raise AssertionError(f"unreachable Di4C student init mode: {mode}")
+
+    load_small_student._dlb_student_init_patch = True  # type: ignore[attr-defined]
+    sdtt.load_small_student = load_small_student
+
+
+def main() -> int:
+    _install_upstream_path()
+    _register_omegaconf_resolvers()
+    import hydra
+
+    _patch_student_loader()
+    import sdtt.main as upstream_main
+
+    @hydra.main(
+        version_base=None,
+        config_path="../upstreams/di4c/sdtt/src/sdtt/configs",
+        config_name="config",
+    )
+    def run(config: Any) -> None:
+        upstream_main.main.__wrapped__(config)
+
+    return int(run() or 0)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

adapters/train_duo_dcd.py

@@ -0,0 +1,58 @@
+"""Run Duo+DCD with a narrow compatibility patch for the pinned Duo source."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import sys
+from typing import Any
+
+
+def _patch_distillation_nll() -> None:
+    import algo
+
+    original = algo.Distillation.nll
+    if getattr(original, "_dlb_accepts_labels", False):
+        return
+
+    def nll(
+        self: Any,
+        x0: Any,
+        labels_or_output_tokens: Any = None,
+        output_tokens: Any = None,
+        current_accumulation_step: Any = None,
+        train_mode: Any = None,
+    ) -> Any:
+        if output_tokens is None:
+            output_tokens = labels_or_output_tokens
+        return original(
+            self,
+            x0,
+            output_tokens,
+            current_accumulation_step,
+            train_mode,
+        )
+
+    nll._dlb_accepts_labels = True  # type: ignore[attr-defined]
+    algo.Distillation.nll = nll
+
+
+def main() -> int:
+    import hydra
+
+    sys.path.insert(0, str(Path.cwd()))
+    _patch_distillation_nll()
+    import main as upstream_main
+
+    @hydra.main(
+        version_base=None,
+        config_path="../upstreams/duo/configs",
+        config_name="config",
+    )
+    def run(config: Any) -> None:
+        upstream_main.main.__wrapped__(config)
+
+    return int(run() or 0)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

artifacts/data.yaml

@@ -0,0 +1,33 @@
+schema_version: 1
+resolved_at: "2026-07-29"
+datasets:
+  lm1b:
+    repo_id: billion-word-benchmark/lm1b
+    revision: 35161838ea9e05371a25a8db037f94fcae4c2064
+    parquet_revision: 8d52bfd3cc2819fc1166dbb8144c328e2690de3e
+    config: plain_text
+    splits:
+      train: train
+      validation: test
+    sequence_length: 128
+    tokenizer: bert-base-uncased
+  owt:
+    repo_id: Skylion007/openwebtext
+    revision: 79d93d786212f7344586290adb811d4ae6a1762c
+    config: plain_text
+    documents: 8013769
+    splits:
+      train: "train[:-100000]"
+      validation: "train[-100000:]"
+    sequence_length: 1024
+    tokenizer: gpt2
+models:
+  bert-base-uncased: 86b5e0934494bd15c9632b12f734a8a67f723594
+  gpt2: 607a30d783dfa663caf39e06633721c8d4cfcd7e
+  gpt2-large: 32b71b12589c2f8d625668d2335a01cac3249519
+lm1b_archive:
+  url: http://www.statmt.org/lm-benchmark/1-billion-word-language-modeling-benchmark-r13output.tar.gz
+  size_bytes: 1792209805
+  sha256: 01ba60381110baf7f189dfd2b8374de371e8c9a340835793f190bdae9e90a34e
+  train_documents: 30301028
+  test_documents: 306688

checkpoints/official/flm/lm1b/model.safetensors

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

checkpoints/official/mdlm/owt/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47149e73f7552f39ea9776dbe74d925d25237bcf2ed2e2ec03cdff9d51c82aa4
+size 678522728

data/elf/assets.json

@@ -0,0 +1,171 @@
+{
+  "assets": {
+    "owt": {
+      "repo_id": "embedded-language-flows/ELF-B-owt-torch",
+      "revision": "146f84133c1389bfd4ef47f14ec7a955da22faa7",
+      "repo_type": "model",
+      "patterns": [
+        "checkpoint_95085",
+        "config.yml"
+      ],
+      "checkpoint": "checkpoint_95085",
+      "path": "data/elf/snapshots/owt/146f84133c1389bfd4ef47f14ec7a955da22faa7",
+      "sha256": {
+        "checkpoint_95085": "52c21bfb7f920e785e0e780645c411c94e4541af4cee8f85b6f555f826d4c283",
+        "config.yml": "cce48467f0a54b99d9512e2d0bc670f9cddcc211436c98a0076591f0be9ebfa8"
+      }
+    },
+    "wmt14": {
+      "repo_id": "embedded-language-flows/ELF-B-de-en-torch",
+      "revision": "93ce98315a2dec985cdefb8d8e62ab618e896ec4",
+      "repo_type": "model",
+      "patterns": [
+        "checkpoint_880600",
+        "config.yml"
+      ],
+      "checkpoint": "checkpoint_880600",
+      "path": "data/elf/snapshots/wmt14/93ce98315a2dec985cdefb8d8e62ab618e896ec4",
+      "sha256": {
+        "checkpoint_880600": "46d77302b9f28b293227a2d29fd0f20b50be289ce9aa63a32f515f3290bbe7cc",
+        "config.yml": "1ac04147a1d80c4f8cd7fd304d11fc3897c490b1dfd132c9d472d6f781dd9657"
+      }
+    },
+    "xsum": {
+      "repo_id": "embedded-language-flows/ELF-B-xsum-torch",
+      "revision": "aca4abd1c7ff5f841f0e2a3ec8a18b0a927927f1",
+      "repo_type": "model",
+      "patterns": [
+        "checkpoint_39800",
+        "config.yml"
+      ],
+      "checkpoint": "checkpoint_39800",
+      "path": "data/elf/snapshots/xsum/aca4abd1c7ff5f841f0e2a3ec8a18b0a927927f1",
+      "sha256": {
+        "checkpoint_39800": "fb2cb9a2aba0328d0de30d53800c0976a019303b9a31f16051981982ed8795ae",
+        "config.yml": "70a3850da7244a5aa6ca61698233945f56ea2c44f5b0c73f3c7fd42375f8cdf3"
+      }
+    },
+    "t5": {
+      "repo_id": "google-t5/t5-small",
+      "revision": "df1b051c49625cf57a3d0d8d3863ed4d13564fe4",
+      "repo_type": "model",
+      "patterns": [
+        "config.json",
+        "model.safetensors",
+        "spiece.model",
+        "tokenizer.json",
+        "tokenizer_config.json"
+      ],
+      "path": "data/elf/snapshots/t5/df1b051c49625cf57a3d0d8d3863ed4d13564fe4",
+      "sha256": {
+        "config.json": "530e25060e3a8d5f7b0fcf53bfea9f3601165161f3b1914676e98d97cf07bcf1",
+        "model.safetensors": "bd944e5f1b3ad9b70dd9d00010a517059e19265671076b8b0a4a58d9491842bc",
+        "spiece.model": "d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86",
+        "tokenizer.json": "d2acde0d8d71dd30a711834b07781b9c89feaac33fd332f60507699282740066",
+        "tokenizer_config.json": "d1e7146101aa96282057f374aa9c3b260fba2109e5990edb1775d6efc70ffa3c"
+      }
+    },
+    "gpt2": {
+      "repo_id": "gpt2",
+      "revision": "607a30d783dfa663caf39e06633721c8d4cfcd7e",
+      "repo_type": "model",
+      "patterns": [
+        "config.json",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "vocab.json",
+        "merges.txt"
+      ],
+      "path": "data/elf/snapshots/gpt2/607a30d783dfa663caf39e06633721c8d4cfcd7e",
+      "sha256": {
+        "config.json": "0daed7749b4f02b8f76240d5444551d7b08712dab4d0adb8239c56ba823bb7b4",
+        "merges.txt": "1ce1664773c50f3e0cc8842619a93edc4624525b728b188a9e0be33b7726adc5",
+        "tokenizer.json": "8414cab924d8b9b33013f0d221c5862f365ee9be39c5c2bfae8a5a9e970478a6",
+        "tokenizer_config.json": "5e04eb606e3a1583530a42e36c2a6b6615c86f34fe77e44d9ddeb43ff940931f",
+        "vocab.json": "196139668be63f3b5d6574427317ae82f612a97c5d1cdaf36ed2256dbf636783"
+      }
+    },
+    "gpt2-large": {
+      "repo_id": "gpt2-large",
+      "revision": "32b71b12589c2f8d625668d2335a01cac3249519",
+      "repo_type": "model",
+      "patterns": [
+        "config.json",
+        "model.safetensors",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "vocab.json",
+        "merges.txt"
+      ],
+      "path": "data/elf/snapshots/gpt2-large/32b71b12589c2f8d625668d2335a01cac3249519",
+      "sha256": {
+        "config.json": "7fccdcfd6622055342a734c663ee0b61ff4fd697f42467595df0bf4448c8c170",
+        "merges.txt": "1ce1664773c50f3e0cc8842619a93edc4624525b728b188a9e0be33b7726adc5",
+        "model.safetensors": "5f47f3e12f91cd33b662ce7e433b6150ad5512b5884a2cee961b50e9c3bbebce",
+        "tokenizer.json": "8414cab924d8b9b33013f0d221c5862f365ee9be39c5c2bfae8a5a9e970478a6",
+        "tokenizer_config.json": "5e04eb606e3a1583530a42e36c2a6b6615c86f34fe77e44d9ddeb43ff940931f",
+        "vocab.json": "196139668be63f3b5d6574427317ae82f612a97c5d1cdaf36ed2256dbf636783"
+      }
+    },
+    "wmt14-official-validation": {
+      "repo_id": "embedded-language-flows/wmt14_de-en_validation_t5",
+      "revision": "bcc92bda4b07c329354843724c9d9ea740f30685",
+      "repo_type": "dataset",
+      "patterns": [
+        "*.arrow",
+        "*.json"
+      ],
+      "path": "data/elf/snapshots/wmt14-official-validation/bcc92bda4b07c329354843724c9d9ea740f30685",
+      "sha256": {
+        "data-00000-of-00001.arrow": "fbd26e8d20fb89d70283450ecda1635fb832f9cb42b713b1ac9c491f464cbd1e",
+        "dataset_info.json": "bf037f9baf98ca5b4fae81e09c0328c449fd28a35235d35500c452d84d8cd377",
+        "state.json": "106bdc18361a8b2b9d2a2a71b1c4139d316a15d824af84aa36140a87ec23e55d"
+      }
+    },
+    "xsum-official-validation": {
+      "repo_id": "embedded-language-flows/xsum_validation_t5",
+      "revision": "b0b4f5de92434d308eceeb9e38fd3f8d7d2b6a9e",
+      "repo_type": "dataset",
+      "patterns": [
+        "*.arrow",
+        "*.json"
+      ],
+      "path": "data/elf/snapshots/xsum-official-validation/b0b4f5de92434d308eceeb9e38fd3f8d7d2b6a9e",
+      "sha256": {
+        "data-00000-of-00001.arrow": "375d55bc9d0872ff172e0f5cff4e8bd1757fb61efd8a1c0d94c05713bdd901d8",
+        "dataset_info.json": "56049acaa63e7257169adeffc460b9d339d349bced668af515cb8cbca0da6b76",
+        "state.json": "1c7b7b52a9e5959b450e41c9ef79069a493386c8990a236207cd507c91631b62"
+      }
+    },
+    "wmt14-raw": {
+      "repo_id": "wmt/wmt14",
+      "revision": "b199e406369ec1b7634206d3ded5ba45de2fe696",
+      "repo_type": "dataset",
+      "patterns": [
+        "de-en/validation-*.parquet",
+        "de-en/test-*.parquet"
+      ],
+      "path": "data/elf/snapshots/wmt14-raw/b199e406369ec1b7634206d3ded5ba45de2fe696",
+      "sha256": {
+        "de-en/test-00000-of-00001.parquet": "a3baab716f14f8abfa63e54995e31f035bb78b20c7943de72f3e1110b7ff35bb",
+        "de-en/validation-00000-of-00001.parquet": "c4a2a0e3b3e472f97f4eceb6d06724571f5390ac705768a02c1a9ad5c3d675ea"
+      }
+    },
+    "xsum-raw": {
+      "repo_id": "EdinburghNLP/xsum",
+      "revision": "7d4d486c2f8ef850b1a11aead99b894ff3dd7da9",
+      "repo_type": "dataset",
+      "patterns": [
+        "data/validation-*.parquet",
+        "data/test-*.parquet"
+      ],
+      "path": "data/elf/snapshots/xsum-raw/7d4d486c2f8ef850b1a11aead99b894ff3dd7da9",
+      "sha256": {
+        "data/test-00000-of-00001.parquet": "224e9dbc6fed987759c1954603b43cb280b8d475d78893779130aa707d967ed7",
+        "data/validation-00000-of-00001.parquet": "b7fc808c2dd3d1340349481aeef489e6eb8c22ff5739d8bce6d0f99e5ef8044d"
+      }
+    }
+  },
+  "source_commit": "b29d8833609e9ab7f67cd9da39435ac5cea04837",
+  "lock_sha256": "b801cd97165cfd78a9f41d6dcbaa4bbf51c60df794237e8a981f8b1520b896a5"
+}

data/elf/owt-prefix.manifest.json

@@ -0,0 +1,8 @@
+{
+  "schema": "elf-input-v1",
+  "count": 1024,
+  "sha256": "3df4bc0c58454494944c86e374388f83ef3a4c2168a6c06f6b70a217cc8e3958",
+  "task": "owt-prefix",
+  "protocol": "c64_text_t5_v1",
+  "source_sha256": "68907c7d4eb5104dbaea6d136a951635a11eb46c0f3e8f417d5f406e1ca834fc"
+}

data/elf/snapshots/gpt2-large/32b71b12589c2f8d625668d2335a01cac3249519/config.json

@@ -0,0 +1,31 @@
+{
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 1280,
+  "n_head": 20,
+  "n_layer": 36,
+  "n_positions": 1024,
+  "resid_pdrop": 0.1,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "vocab_size": 50257
+}

data/elf/snapshots/gpt2-large/32b71b12589c2f8d625668d2335a01cac3249519/tokenizer_config.json

@@ -0,0 +1 @@
+{"model_max_length": 1024}

data/elf/snapshots/gpt2/607a30d783dfa663caf39e06633721c8d4cfcd7e/config.json

@@ -0,0 +1,31 @@
+{
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 768,
+  "n_head": 12,
+  "n_layer": 12,
+  "n_positions": 1024,
+  "resid_pdrop": 0.1,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "vocab_size": 50257
+}

data/elf/snapshots/gpt2/607a30d783dfa663caf39e06633721c8d4cfcd7e/tokenizer_config.json

@@ -0,0 +1 @@
+{"model_max_length": 1024}

data/elf/snapshots/owt/146f84133c1389bfd4ef47f14ec7a955da22faa7/config.yml

@@ -0,0 +1,42 @@
+data_path: embedded-language-flows/openwebtext-t5
+max_length: 1024
+encoder_model_name: t5-small
+latent_mean: 0.0
+latent_std: 0.2
+model: ELF-B
+bottleneck_dim: 128
+num_time_tokens: 4
+num_self_cond_cfg_tokens: 4
+num_model_mode_tokens: 4
+denoiser_p_mean: -1.5
+denoiser_p_std: 0.8
+denoiser_noise_scale: 2.0
+t_eps: 0.05
+time_schedule: logit_normal
+decoder_prob: 0.2
+decoder_noise_scale: 5.0
+decoder_p_mean: 0.8
+decoder_p_std: 0.8
+self_cond_prob: 0.5
+epochs: 5
+global_batch_size: 512
+blr: 0.001
+weight_decay: 0.0
+warmup_epochs: 0.5
+warmup_steps: -1
+optimizer: muon
+ema_decay1: 0.9999
+sampling_configs_path: src/configs/sampling_configs/uncond_sampling_configs.yml
+num_samples: 1000
+log_freq: 100
+save_freq: 1
+save_last_freq: 1
+eval_freq: 1
+output_dir: outputs/elf_b-owt_torch
+resume: null
+use_wandb: true
+wandb_project: elf
+wandb_entity: null
+wandb_run_name: elf_b-owt
+seed: 42
+hf_repo_id: null

data/elf/snapshots/t5/df1b051c49625cf57a3d0d8d3863ed4d13564fe4/config.json

@@ -0,0 +1,51 @@
+{
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 512,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "n_positions": 512,
+  "num_heads": 8,
+  "num_layers": 6,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_num_buckets": 32,
+  "task_specific_params": {
+    "summarization": {
+      "early_stopping": true,
+      "length_penalty": 2.0,
+      "max_length": 200,
+      "min_length": 30,
+      "no_repeat_ngram_size": 3,
+      "num_beams": 4,
+      "prefix": "summarize: "
+    },
+    "translation_en_to_de": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to German: "
+    },
+    "translation_en_to_fr": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to French: "
+    },
+    "translation_en_to_ro": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to Romanian: "
+    }
+  },
+  "vocab_size": 32128
+}

data/elf/snapshots/t5/df1b051c49625cf57a3d0d8d3863ed4d13564fe4/tokenizer_config.json

@@ -0,0 +1,111 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

data/elf/snapshots/wmt14-official-validation/bcc92bda4b07c329354843724c9d9ea740f30685/dataset_info.json

@@ -0,0 +1,101 @@
+{
+  "builder_name": "parquet",
+  "citation": "",
+  "config_name": "de-en",
+  "dataset_name": "wmt14",
+  "dataset_size": 1359920533,
+  "description": "",
+  "download_checksums": {
+    "hf://datasets/wmt/wmt14@b199e406369ec1b7634206d3ded5ba45de2fe696/de-en/train-00000-of-00003.parquet": {
+      "num_bytes": 279527247,
+      "checksum": null
+    },
+    "hf://datasets/wmt/wmt14@b199e406369ec1b7634206d3ded5ba45de2fe696/de-en/train-00001-of-00003.parquet": {
+      "num_bytes": 264970514,
+      "checksum": null
+    },
+    "hf://datasets/wmt/wmt14@b199e406369ec1b7634206d3ded5ba45de2fe696/de-en/train-00002-of-00003.parquet": {
+      "num_bytes": 272987200,
+      "checksum": null
+    },
+    "hf://datasets/wmt/wmt14@b199e406369ec1b7634206d3ded5ba45de2fe696/de-en/validation-00000-of-00001.parquet": {
+      "num_bytes": 473798,
+      "checksum": null
+    },
+    "hf://datasets/wmt/wmt14@b199e406369ec1b7634206d3ded5ba45de2fe696/de-en/test-00000-of-00001.parquet": {
+      "num_bytes": 508753,
+      "checksum": null
+    }
+  },
+  "download_size": 818467512,
+  "features": {
+    "input": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "target": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "condition_input_ids": {
+      "feature": {
+        "dtype": "int64",
+        "_type": "Value"
+      },
+      "_type": "List"
+    },
+    "condition_sequence_length": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "input_ids": {
+      "feature": {
+        "dtype": "int32",
+        "_type": "Value"
+      },
+      "_type": "List"
+    },
+    "sequence_length": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": "",
+  "size_in_bytes": 2178388045,
+  "splits": {
+    "train": {
+      "name": "train",
+      "num_bytes": 1358406800,
+      "num_examples": 4508785,
+      "shard_lengths": [
+        1521929,
+        1721928,
+        1264928
+      ],
+      "dataset_name": "wmt14"
+    },
+    "validation": {
+      "name": "validation",
+      "num_bytes": 736407,
+      "num_examples": 3000,
+      "dataset_name": "wmt14"
+    },
+    "test": {
+      "name": "test",
+      "num_bytes": 777326,
+      "num_examples": 3003,
+      "dataset_name": "wmt14"
+    }
+  },
+  "version": {
+    "version_str": "0.0.0",
+    "major": 0,
+    "minor": 0,
+    "patch": 0
+  }
+}

data/elf/snapshots/wmt14-official-validation/bcc92bda4b07c329354843724c9d9ea740f30685/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "5f50308c846b2747",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

data/elf/snapshots/wmt14/93ce98315a2dec985cdefb8d8e62ab618e896ec4/config.yml

@@ -0,0 +1,45 @@
+data_path: embedded-language-flows/wmt14_de-en_train_t5
+eval_data_path: embedded-language-flows/wmt14_de-en_validation_t5
+max_length: 128
+max_input_length: 64
+pad_token: eos
+encoder_model_name: t5-small
+latent_mean: 0.0
+latent_std: 0.2
+model: ELF-B
+bottleneck_dim: 128
+num_time_tokens: 4
+num_self_cond_cfg_tokens: 4
+num_model_mode_tokens: 4
+denoiser_p_mean: -1.5
+denoiser_p_std: 0.8
+denoiser_noise_scale: 2.0
+t_eps: 0.05
+time_schedule: logit_normal
+decoder_prob: 0.2
+decoder_noise_scale: 5.0
+decoder_p_mean: 0.8
+decoder_p_std: 0.8
+label_drop_prob: 0.1
+self_cond_prob: 0.5
+epochs: 100
+global_batch_size: 512
+blr: 0.001
+weight_decay: 0.0
+warmup_steps: 5000
+optimzier: muon
+ema_decay1: 0.9999
+sampling_configs_path: src/configs/sampling_configs/cond_sampling_configs.yml
+num_samples: 5000
+log_freq: 100
+save_freq: 10
+save_last_freq: 1000
+eval_freq: 10
+output_dir: outputs/elf_b-de-en_from_torch
+resume: null
+use_wandb: true
+wandb_project: elf
+wandb_entity: null
+wandb_run_name: elf_b-wmt-de-en
+seed: 42
+hf_repo_id: null

data/elf/snapshots/xsum-official-validation/b0b4f5de92434d308eceeb9e38fd3f8d7d2b6a9e/dataset_info.json

@@ -0,0 +1,88 @@
+{
+  "builder_name": "parquet",
+  "citation": "",
+  "config_name": "default",
+  "dataset_name": "xsum",
+  "dataset_size": 532255381,
+  "description": "",
+  "download_checksums": {
+    "hf://datasets/EdinburghNLP/xsum@7d4d486c2f8ef850b1a11aead99b894ff3dd7da9/data/train-00000-of-00001.parquet": {
+      "num_bytes": 299612457,
+      "checksum": null
+    },
+    "hf://datasets/EdinburghNLP/xsum@7d4d486c2f8ef850b1a11aead99b894ff3dd7da9/data/validation-00000-of-00001.parquet": {
+      "num_bytes": 16447668,
+      "checksum": null
+    },
+    "hf://datasets/EdinburghNLP/xsum@7d4d486c2f8ef850b1a11aead99b894ff3dd7da9/data/test-00000-of-00001.parquet": {
+      "num_bytes": 16731226,
+      "checksum": null
+    }
+  },
+  "download_size": 332791351,
+  "features": {
+    "input": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "target": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "condition_input_ids": {
+      "feature": {
+        "dtype": "int64",
+        "_type": "Value"
+      },
+      "_type": "List"
+    },
+    "condition_sequence_length": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "input_ids": {
+      "feature": {
+        "dtype": "int32",
+        "_type": "Value"
+      },
+      "_type": "List"
+    },
+    "sequence_length": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": "",
+  "size_in_bytes": 865046732,
+  "splits": {
+    "train": {
+      "name": "train",
+      "num_bytes": 479206363,
+      "num_examples": 204045,
+      "dataset_name": "xsum"
+    },
+    "validation": {
+      "name": "validation",
+      "num_bytes": 26292877,
+      "num_examples": 11332,
+      "dataset_name": "xsum"
+    },
+    "test": {
+      "name": "test",
+      "num_bytes": 26756141,
+      "num_examples": 11334,
+      "dataset_name": "xsum"
+    }
+  },
+  "version": {
+    "version_str": "0.0.0",
+    "major": 0,
+    "minor": 0,
+    "patch": 0
+  }
+}

data/elf/snapshots/xsum-official-validation/b0b4f5de92434d308eceeb9e38fd3f8d7d2b6a9e/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "f934ca16a682895c",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

data/elf/snapshots/xsum/aca4abd1c7ff5f841f0e2a3ec8a18b0a927927f1/config.yml

@@ -0,0 +1,45 @@
+data_path: embedded-language-flows/xsum_train_t5
+eval_data_path: embedded-language-flows/xsum_validation_t5
+max_length: 1088
+max_input_length: 1024
+pad_token: eos
+encoder_model_name: t5-small
+latent_mean: 0.0
+latent_std: 0.2
+model: ELF-B
+bottleneck_dim: 128
+num_time_tokens: 4
+num_self_cond_cfg_tokens: 4
+num_model_mode_tokens: 4
+denoiser_p_mean: -1.5
+denoiser_p_std: 0.8
+denoiser_noise_scale: 2.0
+t_eps: 0.05
+time_schedule: logit_normal
+decoder_prob: 0.2
+decoder_noise_scale: 5.0
+decoder_p_mean: 0.8
+decoder_p_std: 0.8
+label_drop_prob: 0.1
+self_cond_prob: 0.5
+epochs: 100
+global_batch_size: 512
+blr: 0.001
+weight_decay: 0.0
+warmup_steps: 5000
+optimizer: muon
+ema_decay1: 0.9999
+sampling_configs_path: src/configs/sampling_configs/cond_sampling_configs.yml
+num_samples: 5000
+log_freq: 100
+save_freq: 10
+save_last_freq: 1000
+eval_freq: 10
+output_dir: outputs/elf_b-xsum_torch
+resume: null
+use_wandb: true
+wandb_project: elf
+wandb_entity: null
+wandb_run_name: elf_b-xsum
+seed: 42
+hf_repo_id: null

data/elf/wmt14-official-validation.manifest.json

@@ -0,0 +1,17 @@
+{
+  "schema": "elf-input-v1",
+  "count": 3000,
+  "sha256": "1f40387571aa2d3543199a414a7539fcb6a2b6eae18e2d3d3d5765c7480f8e30",
+  "task": "wmt14",
+  "split": "official-validation",
+  "asset": {
+    "repo_id": "embedded-language-flows/wmt14_de-en_validation_t5",
+    "revision": "bcc92bda4b07c329354843724c9d9ea740f30685",
+    "repo_type": "dataset",
+    "patterns": [
+      "*.arrow",
+      "*.json"
+    ]
+  },
+  "reader": "pyarrow_ipc_v1"
+}

data/elf/wmt14-test.manifest.json

@@ -0,0 +1,17 @@
+{
+  "schema": "elf-input-v1",
+  "count": 3003,
+  "sha256": "df5bb91d2d3ab7a2f85cb278c0d70fbf419c58aee8b2435ea488359941a9dc5a",
+  "task": "wmt14",
+  "split": "test",
+  "asset": {
+    "repo_id": "wmt/wmt14",
+    "revision": "b199e406369ec1b7634206d3ded5ba45de2fe696",
+    "repo_type": "dataset",
+    "patterns": [
+      "de-en/validation-*.parquet",
+      "de-en/test-*.parquet"
+    ]
+  },
+  "reader": "pyarrow_parquet_v1"
+}

data/elf/wmt14-validation.manifest.json

@@ -0,0 +1,17 @@
+{
+  "schema": "elf-input-v1",
+  "count": 3000,
+  "sha256": "ec0b7b23c62af5c941d2d6d6eda75588d0e1d99e13a68e580ac0f37583687680",
+  "task": "wmt14",
+  "split": "validation",
+  "asset": {
+    "repo_id": "wmt/wmt14",
+    "revision": "b199e406369ec1b7634206d3ded5ba45de2fe696",
+    "repo_type": "dataset",
+    "patterns": [
+      "de-en/validation-*.parquet",
+      "de-en/test-*.parquet"
+    ]
+  },
+  "reader": "pyarrow_parquet_v1"
+}