logits-cookbook 0.1.0__tar.gz

logits_cookbook-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,49 @@
+name: CI
+
+on:
+  push:
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  lint:
+    name: Lint (ruff)
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install ruff
+        run: python -m pip install --upgrade pip ruff
+
+      - name: ruff check
+        run: ruff check .
+
+      - name: ruff format --check
+        run: ruff format --check .
+
+  package:
+    name: Build package
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Build and check distributions
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install build twine
+          python -m build
+          python -m twine check dist/*

logits_cookbook-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,87 @@
+name: Publish
+
+on:
+  release:
+    types: [published]
+  workflow_dispatch:
+    inputs:
+      repository:
+        description: "Repository to publish to"
+        required: true
+        default: "testpypi"
+        type: choice
+        options:
+          - testpypi
+          - pypi
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    name: Build distributions
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Build package
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install build twine
+          python -m build
+          python -m twine check dist/*
+
+      - name: Upload distributions
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+
+  publish-testpypi:
+    name: Publish to TestPyPI
+    needs: build
+    if: github.event_name == 'workflow_dispatch' && inputs.repository == 'testpypi'
+    runs-on: ubuntu-latest
+    environment: testpypi
+    permissions:
+      id-token: write
+      contents: read
+
+    steps:
+      - name: Download distributions
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+
+      - name: Publish distributions to TestPyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+  publish-pypi:
+    name: Publish to PyPI
+    needs: build
+    if: github.event_name == 'release' || (github.event_name == 'workflow_dispatch' && inputs.repository == 'pypi')
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+      contents: read
+
+    steps:
+      - name: Download distributions
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+
+      - name: Publish distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

logits_cookbook-0.1.0/.gitignore

@@ -0,0 +1,8 @@
+**/__pycache__
+.DS_Store
+.env
+.env.*
+.venv
+uv.lock
+sdk
+logits-sdk

logits_cookbook-0.1.0/AGENTS.md

@@ -0,0 +1,135 @@
+# Logits Cookbook Agent Guide
+
+Quick reference for agents working on `logits_cookbook`. Full documentation is in `docs/`.
+
+`logits_cookbook` is a client library with training and eval code built on the Tinker service (hosted by Thinking Machines Lab) and the Tinker SDK (a separate repo with just the API). You author training/eval loops that run on a CPU machine; Tinker executes the heavy GPU work.
+
+**Start here:** `docs/training-sampling.mdx` - Complete walkthrough of training and sampling basics.
+
+## Documentation Map (`docs/`)
+
+**API Fundamentals:**
+- `index.mdx` - Tinker overview, division of responsibilities
+- `install.mdx` - Installation, API key setup
+- `training-sampling.mdx` - **Starter guide**: data prep, forward_backward, sampling, vision inputs
+- `losses.mdx` - Loss functions (cross_entropy, importance_sampling, ppo, cispo, dro, forward_backward_custom)
+- `save-load.mdx` - Checkpointing (save_weights_for_sampler vs save_state)
+- `async.mdx` - Sync/async APIs, futures, overlapping requests
+- `model-lineup.mdx` - Available models
+- `under-the-hood.mdx` - Clock cycles, worker pools
+
+**API Reference (`api-reference/`):**
+- `types.md` - **All API types** (Datum, ModelInput, TensorData, SamplingParams, etc.)
+- `trainingclient.md`, `samplingclient.md`, `serviceclient.md`, `restclient.md` - Client APIs
+
+**Supervised Learning (`supervised-learning/`):**
+- `../supervised-learning.mdx` - SL overview
+- `sl-basic.mdx` - First SL run
+- `sl-hyperparams.mdx` - LR formula, batch size
+- `sl-loop.mdx` - Minimal training loop
+- `prompt-distillation.mdx` - Distilling prompts
+- `sweep-case-study.mdx` - Hyperparameter sweeps
+
+**Reinforcement Learning (`rl/`):**
+- `../rl.mdx` - RL overview (RLVR, RLHF)
+- `rl-basic.mdx` - First RL run
+- `rl-envs.mdx` - Custom Env, EnvGroupBuilder, RLDataset
+- `rl-loops.mdx` - Minimal RL loop
+- `rl-hyperparams.mdx` - batch_size vs group_size, async training
+- `sequence-extension.mdx` - Multi-turn RL, KV-cache
+
+**Preferences (`preferences/`):**
+- `../preferences.mdx` - DPO vs RLHF overview
+- `dpo-guide.mdx` - DPO training
+- `rlhf-example.mdx` - RLHF pipeline
+
+**Other:**
+- `rendering.mdx` - Renderers (bridge between chat-style data and token sequences), vision inputs, TrainOnWhat
+- `completers.mdx` - TokenCompleter vs MessageCompleter
+- `evals.mdx` - Inline evals, Inspect AI, custom evaluators
+- `lora-primer.mdx` - LoRA background
+- `download-weights.mdx` / `publish-weights.mdx` - Weight export
+
+---
+
+## Composing Types
+
+Agents often struggle with the nested type hierarchy. Key resources:
+
+**Reference:** `docs/api-reference/types.md` documents all API types.
+
+**Core types:**
+- `Datum` = `model_input` (ModelInput) + `loss_fn_inputs` (dict of TensorData)
+- `ModelInput` = list of chunks (EncodedTextChunk, ImageChunk)
+- `TensorData` = wrapper for numpy/torch arrays with shape info
+
+**Helper functions** (use these instead of manual construction):
+- `datum_from_model_input_weights(model_input, weights, max_length)` - SL datum creation (`supervised/common.py`)
+- `conversation_to_datum(messages, renderer, max_length, train_on_what)` - Full pipeline (`supervised/data.py`)
+- `renderer.build_supervised_example(messages)` - Returns (ModelInput, weights)
+- `ModelInput.from_ints(tokens)` - Create from token list
+- `TensorData.from_numpy(arr)` / `TensorData.from_torch(tensor)` - Wrap arrays
+
+---
+
+## Architecture
+
+**Builder pattern:** Config objects are `chz` dataclasses (SupervisedDatasetBuilder, RLDatasetBuilder, EnvGroupBuilder). They expose `.build()`/`__call__()` returning runtime objects.
+
+**Key code locations:**
+- SL: `logits_cookbook/supervised/train.py`
+- RL: `logits_cookbook/rl/train.py`
+- DPO: `logits_cookbook/preference/train_dpo.py`
+- Renderers: `logits_cookbook/renderers/`
+- Completers: `logits_cookbook/completers.py`
+- RL types: `logits_cookbook/rl/types.py`
+- Logging: `logits_cookbook/utils/logtree.py`, `logits_cookbook/rl/rollouts.py`
+- Recipes: `logits_cookbook/recipes/`
+
+**Training outputs:** RL and SL training write human-readable HTML reports and machine-readable JSON files (metrics, rollout transcripts, per-trajectory summaries) to `log_path`. Point agents at a `log_path` directory to analyze training runs — `metrics.jsonl` for scalar metrics, `*_rollout_summaries.jsonl` for per-trajectory data, and `*_logtree.json` for full rollout transcripts including model responses. See `docs/rl/rl-logging.mdx` for the complete file reference and parsing examples.
+
+---
+
+## Conventions
+
+**Subscript suffixes** for tensor names: `_P` (problems), `_G` (groups), `_T` (tokens), `_D` (datums). Example: `tokens_P_G_T[p][g][t]`
+
+**Code style:**
+- Explicit typing; avoid `Any` / `type: ignore`
+- Use `safezip`, `timed`, `scope` helpers
+- `@chz.chz` decorator for config serialization
+- `ml_log.log_metrics` for metrics; `logtree` for transcripts
+
+**Env lifecycle:** `Env` objects are single-use (no reset). Create via `EnvGroupBuilder`.
+
+---
+
+## Common Pitfalls
+
+1. **LoRA LR:** Use `hyperparam_utils.get_lr(model_name)` - LoRA needs ~10x higher LR than full fine-tuning.
+
+2. **Renderer mismatch:** Match `renderer_name` to model family (`llama3`, `qwen3`, `role_colon`).
+
+3. **Async gaps:** Submit `forward_backward_async` and `optim_step_async` back-to-back before awaiting.
+
+4. **Sampler desync:** Create a **new** sampling client after saving weights.
+
+5. **Type construction:** Use helper functions, not manual dict construction. See `supervised/data.py` and `supervised/common.py`.
+
+6. **Group semantics:** RL advantages are centered within each group.
+
+7. **DPO:** Start with `dpo_beta=0.1`, LR~1e-5.
+
+---
+
+## Testing
+
+```bash
+# Unit tests (no API needed, colocated *_test.py files)
+pytest logits_cookbook/
+
+# Smoke tests (requires TINKER_API_KEY + network)
+pytest tests/
+```
+
+For debugging, shrink workloads via `n_batches`, `batch_size`, `group_size` in dataset builders.

logits_cookbook-0.1.0/CHANGELOG.md

@@ -0,0 +1,5 @@
+### Fork from tinker cookbook
+**Date:** 2026-04-09
+**Type:** new
+**Tags:** all
+

logits_cookbook-0.1.0/CLAUDE.md

@@ -0,0 +1,135 @@
+# Logits Cookbook Agent Guide
+
+Quick reference for agents working on `logits_cookbook`. Full documentation is in `docs/`.
+
+`logits_cookbook` is a client library with training and eval code built on the Tinker service (hosted by Thinking Machines Lab) and the Tinker SDK (a separate repo with just the API). You author training/eval loops that run on a CPU machine; Tinker executes the heavy GPU work.
+
+**Start here:** `docs/training-sampling.mdx` - Complete walkthrough of training and sampling basics.
+
+## Documentation Map (`docs/`)
+
+**API Fundamentals:**
+- `index.mdx` - Tinker overview, division of responsibilities
+- `install.mdx` - Installation, API key setup
+- `training-sampling.mdx` - **Starter guide**: data prep, forward_backward, sampling, vision inputs
+- `losses.mdx` - Loss functions (cross_entropy, importance_sampling, ppo, cispo, dro, forward_backward_custom)
+- `save-load.mdx` - Checkpointing (save_weights_for_sampler vs save_state)
+- `async.mdx` - Sync/async APIs, futures, overlapping requests
+- `model-lineup.mdx` - Available models
+- `under-the-hood.mdx` - Clock cycles, worker pools
+
+**API Reference (`api-reference/`):**
+- `types.md` - **All API types** (Datum, ModelInput, TensorData, SamplingParams, etc.)
+- `trainingclient.md`, `samplingclient.md`, `serviceclient.md`, `restclient.md` - Client APIs
+
+**Supervised Learning (`supervised-learning/`):**
+- `../supervised-learning.mdx` - SL overview
+- `sl-basic.mdx` - First SL run
+- `sl-hyperparams.mdx` - LR formula, batch size
+- `sl-loop.mdx` - Minimal training loop
+- `prompt-distillation.mdx` - Distilling prompts
+- `sweep-case-study.mdx` - Hyperparameter sweeps
+
+**Reinforcement Learning (`rl/`):**
+- `../rl.mdx` - RL overview (RLVR, RLHF)
+- `rl-basic.mdx` - First RL run
+- `rl-envs.mdx` - Custom Env, EnvGroupBuilder, RLDataset
+- `rl-loops.mdx` - Minimal RL loop
+- `rl-hyperparams.mdx` - batch_size vs group_size, async training
+- `sequence-extension.mdx` - Multi-turn RL, KV-cache
+
+**Preferences (`preferences/`):**
+- `../preferences.mdx` - DPO vs RLHF overview
+- `dpo-guide.mdx` - DPO training
+- `rlhf-example.mdx` - RLHF pipeline
+
+**Other:**
+- `rendering.mdx` - Renderers (bridge between chat-style data and token sequences), vision inputs, TrainOnWhat
+- `completers.mdx` - TokenCompleter vs MessageCompleter
+- `evals.mdx` - Inline evals, Inspect AI, custom evaluators
+- `lora-primer.mdx` - LoRA background
+- `download-weights.mdx` / `publish-weights.mdx` - Weight export
+
+---
+
+## Composing Types
+
+Agents often struggle with the nested type hierarchy. Key resources:
+
+**Reference:** `docs/api-reference/types.md` documents all API types.
+
+**Core types:**
+- `Datum` = `model_input` (ModelInput) + `loss_fn_inputs` (dict of TensorData)
+- `ModelInput` = list of chunks (EncodedTextChunk, ImageChunk)
+- `TensorData` = wrapper for numpy/torch arrays with shape info
+
+**Helper functions** (use these instead of manual construction):
+- `datum_from_model_input_weights(model_input, weights, max_length)` - SL datum creation (`supervised/common.py`)
+- `conversation_to_datum(messages, renderer, max_length, train_on_what)` - Full pipeline (`supervised/data.py`)
+- `renderer.build_supervised_example(messages)` - Returns (ModelInput, weights)
+- `ModelInput.from_ints(tokens)` - Create from token list
+- `TensorData.from_numpy(arr)` / `TensorData.from_torch(tensor)` - Wrap arrays
+
+---
+
+## Architecture
+
+**Builder pattern:** Config objects are `chz` dataclasses (SupervisedDatasetBuilder, RLDatasetBuilder, EnvGroupBuilder). They expose `.build()`/`__call__()` returning runtime objects.
+
+**Key code locations:**
+- SL: `logits_cookbook/supervised/train.py`
+- RL: `logits_cookbook/rl/train.py`
+- DPO: `logits_cookbook/preference/train_dpo.py`
+- Renderers: `logits_cookbook/renderers/`
+- Completers: `logits_cookbook/completers.py`
+- RL types: `logits_cookbook/rl/types.py`
+- Logging: `logits_cookbook/utils/logtree.py`, `logits_cookbook/rl/rollouts.py`
+- Recipes: `logits_cookbook/recipes/`
+
+**Training outputs:** RL and SL training write human-readable HTML reports and machine-readable JSON files (metrics, rollout transcripts, per-trajectory summaries) to `log_path`. Point agents at a `log_path` directory to analyze training runs — `metrics.jsonl` for scalar metrics, `*_rollout_summaries.jsonl` for per-trajectory data, and `*_logtree.json` for full rollout transcripts including model responses. See `docs/rl/rl-logging.mdx` for the complete file reference and parsing examples.
+
+---
+
+## Conventions
+
+**Subscript suffixes** for tensor names: `_P` (problems), `_G` (groups), `_T` (tokens), `_D` (datums). Example: `tokens_P_G_T[p][g][t]`
+
+**Code style:**
+- Explicit typing; avoid `Any` / `type: ignore`
+- Use `safezip`, `timed`, `scope` helpers
+- `@chz.chz` decorator for config serialization
+- `ml_log.log_metrics` for metrics; `logtree` for transcripts
+
+**Env lifecycle:** `Env` objects are single-use (no reset). Create via `EnvGroupBuilder`.
+
+---
+
+## Common Pitfalls
+
+1. **LoRA LR:** Use `hyperparam_utils.get_lr(model_name)` - LoRA needs ~10x higher LR than full fine-tuning.
+
+2. **Renderer mismatch:** Match `renderer_name` to model family (`llama3`, `qwen3`, `role_colon`).
+
+3. **Async gaps:** Submit `forward_backward_async` and `optim_step_async` back-to-back before awaiting.
+
+4. **Sampler desync:** Create a **new** sampling client after saving weights.
+
+5. **Type construction:** Use helper functions, not manual dict construction. See `supervised/data.py` and `supervised/common.py`.
+
+6. **Group semantics:** RL advantages are centered within each group.
+
+7. **DPO:** Start with `dpo_beta=0.1`, LR~1e-5.
+
+---
+
+## Testing
+
+```bash
+# Unit tests (no API needed, colocated *_test.py files)
+pytest logits_cookbook/
+
+# Smoke tests (requires TINKER_API_KEY + network)
+pytest tests/
+```
+
+For debugging, shrink workloads via `n_batches`, `batch_size`, `group_size` in dataset builders.

logits_cookbook-0.1.0/CONTRIBUTING.md

@@ -0,0 +1,77 @@
+# Development
+
+This project is built in the spirit of open science and collaborative development. We believe that the best tools emerge through community involvement and shared learning.
+
+We welcome PR contributions after our private beta is over. If you have any feedback, please email us at tinker@thinkingmachines.ai.
+
+## Organization of training scripts
+
+We're designing the codebase with the following goals:
+
+1. Low barrier to entry: it should be dead simple to run something and see numbers go up.
+2. Extensible: it should be possible to pass in custom datasets and evals and control all the hyperparameters.
+3. Science-friendly: it should be easy to run sweeps, and analyze the results.
+
+To achieve this, we'll use the following structure around training scripts:
+
+- There's a main training function, such as [rl/train.py](logits_cookbook/rl/train.py) or [supervised/train.py](logits_cookbook/supervised/train.py), which contains the main loop.
+    - This function contains a detailed config object (`Config`), which isn't constructable from the command line.
+    - The config contains members that specify things like datasets and evals. However, these should be chz configs (with a `.build` method that constructs the actual object) or callables (we recommend using functools.partial). This way, the config is serializable, which is useful for sweeps.
+- There are launch scripts that assemble training configs (e.g., [recipes/math_rl/train.py](logits_cookbook/recipes/math_rl/train.py)), which construct a smaller config object (`CLIConfig`) from the command line.
+
+## Async
+
+Async is very useful for RL, where it allows us to make many queries in parallel (e.g., sampling calls). For all of the interfaces used in RL (such as the `Env` class), all the methods that take nontrivial amounts of time should be async. For some of the other code, such as [recipes/sl_loop.py](logits_cookbook/recipes/sl_loop.py), we've chosen not to use async methods, just to make it more beginner-friendly, as many python programmers are not familiar with async.
+
+## Typing
+
+Please use typing wherever possible; avoid `Any` and `type: ignore`; prefer casting. However, avoid using convoluted generics or writing code that's much more verbose just to satisfy the type checker. Prefer using single types over union types.
+
+## Classes
+
+There are a lot of different classes, which might make the code feel less approachable. However, they follow *the builder pattern*, and the code should be less confusing when you know the pattern.
+
+We can illustrate the pattern with the two main examples:
+
+- A `SupervisedDatasetBuilder` is a configuration object which builds a `SupervisedDataset`.
+- An `RLDatasetBuilder` is a configuration object which builds an `RLDataset`, which generates batches of `EnvGroupBuilder` objects, which each generate a group of `Env` objects.
+
+Here, the `SupervisedDatasetBuilder`, `RLDatasetBuilder`, and `EnvGroupBuilder` are all configuration objects, which have a `__call__` method that builds another object. You can see these objects in [supervised/types.py](logits_cookbook/supervised/types.py) and [rl/types.py](logits_cookbook/rl/types.py).
+
+In general, we use a lot of configuration objects, with a `__call__` method that returns a heavyweight object (like a dataset). We use `chz` for the configuration objects -- it's similar to a dataclass but with some extra features that are nice for configs. We use either dataclasses or regular python classes for the heavyweight objects.
+
+## Envs
+
+An `Env` is an RL environment. For those with an RL background, it roughly corresponds to an MDP or a POMDP, however we use in more general cases (such as multi-agent settings) that don't strictly correspond to the MDP/POMDP formalism. It's roughly analogous the concept of an Env in OpenAI Gym, but unlike OpenAI Gym, we don't have a `reset` method; rather, the env should be discarded after a rollout. Any shared resources should be maintained by whatever object is creating the envs.
+
+The `Env`s are created by `EnvGroupBuilder`s. The group of envs returned by `EnvGroupBuilder` have something in common; either they correspond to the same task (in which case we can use this information for variance reduction, as in GRPO, which centers per group); or, we can use the group to define a multi-agent environment.
+
+- One common multi-agent environment is where we use a pairwise preference model to compare pairs of completions.
+- We can also use the group to define a two-player game. Some two player games such as tic-tac-toe are currently supported through the [text_arena](logits_cookbook/recipes/multiplayer_rl/text_arena/env.py) environments.
+
+
+## Notation
+
+We'll use subscripts to indicate the shapes of objects. For example, `tokens_P_G_T` indicates a three-dimensional array of tokens, with `P` problems, `G` groups, and `T` tokens per groups, so `tokens_P_G_T[p][g][t]` should refer to a single token. In many cases, the arrays will be ragged. E.g., the `T` axis will have different lengths for different `(p,g)`. Sometimes, a given dimension will be flattened from two dimensions. If we write `tokens_PG_T`, that means that we have a two dimensional array, where the 0th dimension is flattened from the `P` and `G` dimensions.
+
+### Common Dimension Names
+
+Here are the standard dimension subscripts used throughout the codebase:
+
+- `_D`: Data/Datum dimension (for training data items)
+- `_G`: Group dimension (for multiple attempts/rollouts of the same problem)
+- `_P`: Problem dimension (for different problems/prompts)
+- `_T`: Token/Time dimension (for sequences)
+
+The relationship between dimensions in RL:
+- A batch contains multiple problems (`_P`)
+- Each problem spawns multiple attempts/environments (`_G`), forming a group
+- Each attempt produces one trajectory
+- Advantages are normalized within each group (across the `_G` dimension)
+
+Examples:
+- `env_group_builders_P`: A list of environment builders, one per problem
+- `trajectories_G`: Multiple trajectories from attempts at the same problem
+- `rewards_G`: Rewards for each attempt within a group
+- `tokens_P_G_T`: Tokens with problem, group, and time dimensions
+- `data_D`: A list of training data items