Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
6.6 kB · 182 lines
Python
at letta-code
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183#!/usr/bin/env python3"""Bootstrap Qwen tool use with a small audited synthetic SFT corpus."""from __future__ import annotations
import asyncioimport jsonimport mathimport randomimport sysfrom pathlib import Pathfrom typing import Any
import chz
from tinker_cookbook import cli_utils, renderersfrom tinker_cookbook.renderers import ToolCall, TrainOnWhatfrom tinker_cookbook.supervised import trainfrom tinker_cookbook.supervised.data import conversation_to_datumfrom tinker_cookbook.supervised.types import SupervisedDataset, SupervisedDatasetBuilderfrom tinker_cookbook.tokenizer_utils import get_tokenizer
EXAMPLE_DIR = Path(__file__).resolve().parent
def typed_messages(row: dict[str, Any], renderer: renderers.Renderer) -> list[renderers.Message]: prefix = renderer.create_conversation_prefix_with_tools( row.get("tools", []), str(row.get("system_prompt", "")), ) messages: list[renderers.Message] = list(prefix) for raw_message in row["messages"]: message = dict(raw_message) if message.get("tool_calls"): message["tool_calls"] = [ToolCall.model_validate(call) for call in message["tool_calls"]] messages.append(message) # type: ignore[arg-type] return messages
def assistant_training_conversations( row: dict[str, Any], renderer: renderers.Renderer) -> list[list[renderers.Message]]: """Return one full-prefix example for each assistant turn.""" messages = typed_messages(row, renderer) return [messages[: index + 1] for index, message in enumerate(messages) if message["role"] == "assistant"]
class ToolConversationDataset(SupervisedDataset): def __init__( self, rows: list[dict[str, Any]], renderer: renderers.Renderer, batch_size: int, max_length: int, ) -> None: self.original_conversations = [ conversation for row in rows for conversation in assistant_training_conversations(row, renderer) ] self.conversations = list(self.original_conversations) self.renderer = renderer self.batch_size = batch_size self.max_length = max_length
def get_batch(self, index: int): selected = self.conversations[index * self.batch_size : (index + 1) * self.batch_size] return [ conversation_to_datum( conversation, self.renderer, self.max_length, train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE, ) for conversation in selected ]
def set_epoch(self, seed: int = 0) -> None: self.conversations = list(self.original_conversations) random.Random(seed).shuffle(self.conversations)
def __len__(self) -> int: return math.ceil(len(self.conversations) / self.batch_size)
@chz.chzclass ToolDatasetBuilder(SupervisedDatasetBuilder): data_file: str model_name_for_tokenizer: str renderer_name: str batch_size: int = 12 max_length: int = 8192 test_size: int = 3 shuffle_seed: int = 0 language_multipliers: str = ""
def __call__(self) -> tuple[SupervisedDataset, SupervisedDataset | None]: rows = [json.loads(line) for line in Path(self.data_file).read_text(encoding="utf-8").splitlines() if line] if len(rows) <= self.test_size: raise ValueError("tool bootstrap corpus must be larger than the test split") random.Random(self.shuffle_seed).shuffle(rows) encoder = get_tokenizer(self.model_name_for_tokenizer) renderer = renderers.get_renderer(self.renderer_name, encoder) test_rows = rows[: self.test_size] train_rows = rows[self.test_size :] multipliers: dict[str, int] = {} for item in self.language_multipliers.split(","): if not item.strip(): continue language, count = item.split(":", 1) multipliers[language.strip()] = max(1, int(count)) weighted_train_rows = [ row for row in train_rows for _ in range(multipliers.get(str(row.get("metadata", {}).get("language", "")), 1)) ] train_dataset = ToolConversationDataset(weighted_train_rows, renderer, self.batch_size, self.max_length) test_dataset = ToolConversationDataset(test_rows, renderer, self.batch_size, self.max_length) if len(train_dataset) == 0 or len(test_dataset) == 0: raise ValueError("batch size leaves an empty train or test dataset") return train_dataset, test_dataset
@chz.chzclass CLIConfig: data_file: str = str(EXAMPLE_DIR / "data" / "tool-bootstrap.jsonl") model_name: str = "Qwen/Qwen3.5-4B" renderer_name: str = "qwen3_5_disable_thinking" load_checkpoint_path: str | None = None log_path: str = str(EXAMPLE_DIR / "runs" / "tool-bootstrap-sft") batch_size: int = 12 max_length: int = 8192 test_size: int = 3 language_multipliers: str = "" learning_rate: float = 1e-4 num_epochs: int = 2 lora_rank: int = 32 max_steps: int | None = None save_every: int = 9 eval_every: int = 9 behavior_if_log_dir_exists: cli_utils.LogdirBehavior = "ask"
def build_config(cli: CLIConfig) -> train.Config: builder_options = { "data_file": cli.data_file, "model_name_for_tokenizer": cli.model_name, "renderer_name": cli.renderer_name, "batch_size": cli.batch_size, "max_length": cli.max_length, "test_size": cli.test_size, "language_multipliers": cli.language_multipliers, } dataset = ToolDatasetBuilder(**builder_options) return train.Config( log_path=cli.log_path, model_name=cli.model_name, recipe_name="project_euler_tool_bootstrap", load_checkpoint_path=cli.load_checkpoint_path, renderer_name=cli.renderer_name, dataset_builder=dataset, learning_rate=cli.learning_rate, lr_schedule="linear", num_epochs=cli.num_epochs, lora_rank=cli.lora_rank, save_every=cli.save_every, eval_every=cli.eval_every, max_steps=cli.max_steps, )
def main(cli: CLIConfig) -> None: config = build_config(cli) print("Resolved tool bootstrap SFT config:") print(config) cli_utils.check_log_dir(config.log_path, behavior_if_exists=cli.behavior_if_log_dir_exists) asyncio.run(train.main(config))
if __name__ == "__main__": blueprint = chz.Blueprint(CLIConfig) blueprint.make_from_argv(sys.argv[1:]) main(blueprint.make())