| """Independent JSON text SFT baseline; assistant token CE only. |
| |
| Example offline smoke (not a Qwen training result):: |
| |
| python -m fastfill.v2.text_sft --data samples.jsonl --output out/text-smoke \ |
| --backbone tiny --dry-run |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| from pathlib import Path |
| from types import SimpleNamespace |
| from typing import Any |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from .batch import TinyTokenizer, _geometry_rows, load_tokenizer, room_normalization, tokenize_condition |
| from .schema import validate_condition |
| from .io import fingerprint, read_samples, run_metadata, safe_output |
|
|
|
|
| def _answer(sample: dict) -> dict: |
| condition = sample["condition"] |
| validate_condition(condition) |
| origin, scale = room_normalization(condition.get("room", {})) |
| rows = _geometry_rows(sample, origin, scale) |
| if not all(all(row) for row in rows["position_valid"] + rows["size_valid"]) or not all(rows["yaw_valid"]): |
| raise ValueError("text SFT requires complete audited geometry labels") |
| targets = {o["id"]: o for o in sample.get("target", {}).get("objects", [])} |
| return {"schema_version": "fastfill.v2", "objects": [{ |
| "id": obj["id"], "target_size_local_m": targets[obj["id"]]["target_size_local_m"], |
| "bottom_center_m": targets[obj["id"]]["bottom_center_m"], "yaw_rad": targets[obj["id"]]["yaw_rad"], |
| } for obj in condition["objects"]]} |
|
|
|
|
| def collate_text_samples(samples: list[dict], tokenizer: Any, *, max_length: int = 4096) -> dict: |
| if not samples: |
| raise ValueError("cannot batch zero samples") |
| sequences, labels, prompt_lengths = [], [], [] |
| for sample in samples: |
| prompt, _ = tokenize_condition(sample["condition"], tokenizer) |
| answer = json.dumps(_answer(sample), ensure_ascii=False, sort_keys=True, |
| separators=(",", ":"), allow_nan=False) |
| response = tokenizer.encode(answer, add_special_tokens=False) |
| if tokenizer.eos_token_id is not None: |
| response = response + [tokenizer.eos_token_id] |
| if len(prompt) + len(response) > max_length: |
| raise ValueError("text sample exceeds context budget; no partial target truncation") |
| sequences.append(prompt + response) |
| labels.append([-100] * len(prompt) + response) |
| prompt_lengths.append(len(prompt)) |
| length = max(map(len, sequences)) |
| ids = torch.full((len(samples), length), tokenizer.pad_token_id, dtype=torch.long) |
| attention = torch.zeros_like(ids, dtype=torch.bool) |
| label_tensor = torch.full_like(ids, -100) |
| for row, (sequence, target) in enumerate(zip(sequences, labels)): |
| ids[row, :len(sequence)] = torch.tensor(sequence) |
| attention[row, :len(sequence)] = True |
| label_tensor[row, :len(target)] = torch.tensor(target) |
| return {"input_ids": ids, "attention_mask": attention, "labels": label_tensor, |
| "prompt_lengths": prompt_lengths} |
|
|
|
|
| class TinyTextModel(nn.Module): |
| """Causal byte GRU, reserved for baseline plumbing smoke tests.""" |
|
|
| def __init__(self, hidden_size: int = 64): |
| super().__init__() |
| self.embedding = nn.Embedding(258, hidden_size, padding_idx=0) |
| self.encoder = nn.GRU(hidden_size, hidden_size, batch_first=True) |
| self.lm_head = nn.Linear(hidden_size, 258) |
|
|
| def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, **_: Any) -> Any: |
| hidden, _ = self.encoder(self.embedding(input_ids)) |
| return SimpleNamespace(logits=self.lm_head(hidden)) |
|
|
|
|
| def build_text_model(backbone: str, *, hidden_size: int = 64, lora_rank: int = 8, |
| lora_alpha: int = 16, local_files_only: bool = False, |
| backbone_dtype: str = "float32") -> nn.Module: |
| if backbone == "tiny": |
| return TinyTextModel(hidden_size) |
| from transformers import AutoConfig, AutoModelForCausalLM |
| config = AutoConfig.from_pretrained(backbone, local_files_only=local_files_only) |
| if not str(config.model_type).startswith("qwen"): |
| raise ValueError("text baseline backbone must be a Qwen-family checkpoint") |
| model = AutoModelForCausalLM.from_pretrained(backbone, config=config, |
| torch_dtype=getattr(torch, backbone_dtype), local_files_only=local_files_only) |
| model.config.use_cache = False |
| if lora_rank: |
| from peft import LoraConfig, TaskType, get_peft_model |
| model = get_peft_model(model, LoraConfig(task_type=TaskType.CAUSAL_LM, r=lora_rank, |
| lora_alpha=lora_alpha, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"])) |
| return model |
|
|
|
|
| def text_loss(model: nn.Module, batch: dict) -> torch.Tensor: |
| logits = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"]).logits |
| labels = batch["labels"][:, 1:].reshape(-1) |
| if not (labels != -100).any(): |
| raise ValueError("no assistant target tokens available for CE") |
| return F.cross_entropy(logits[:, :-1].float().reshape(-1, logits.shape[-1]), labels, ignore_index=-100) |
|
|
|
|
| def _read_samples(path: Path) -> list[dict]: |
| samples = read_samples(path, training=True) |
| for sample in samples: |
| _answer(sample) |
| return samples |
|
|
|
|
| def _verify_source_fingerprint(metadata): |
| try: |
| matches = fingerprint(metadata["data_path"]) == metadata["data_sha256"] |
| except OSError: |
| matches = False |
| if not matches: |
| raise RuntimeError("text training input fingerprint changed during the run") |
|
|
|
|
| @torch.no_grad() |
| def generate_text(model: nn.Module, tokenizer: Any, condition: dict, *, max_new_tokens: int = 1024, |
| max_length: int = 4096, device: str | torch.device = "cpu") -> str: |
| validate_condition(condition) |
| tokens, _ = tokenize_condition(condition, tokenizer) |
| if len(tokens) + max_new_tokens > max_length: |
| raise ValueError("text generation exceeds configured context budget") |
| model.eval() |
| ids = torch.tensor([tokens], dtype=torch.long, device=device) |
| if not isinstance(model, TinyTextModel): |
| output = model.generate(input_ids=ids, attention_mask=torch.ones_like(ids), |
| max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tokenizer.pad_token_id, |
| eos_token_id=tokenizer.eos_token_id) |
| return tokenizer.decode(output[0, len(tokens):], skip_special_tokens=True) |
| generated = [] |
| hidden, state = model.encoder(model.embedding(ids)) |
| for _ in range(max_new_tokens): |
| next_token = model.lm_head(hidden[:, -1]).argmax(-1) |
| if int(next_token[0]) == tokenizer.eos_token_id: |
| break |
| generated.append(int(next_token[0])) |
| hidden, state = model.encoder(model.embedding(next_token[:, None]), state) |
| return tokenizer.decode(generated) |
|
|
|
|
| def load_text_model(directory: str | Path, *, device: str | torch.device = "cpu", |
| local_files_only: bool = False) -> tuple[nn.Module, Any]: |
| directory = Path(directory) |
| config = json.loads((directory / "text_config.json").read_text()) |
| if config["backbone"] == "tiny": |
| model = build_text_model("tiny", hidden_size=config["hidden_size"]) |
| model.load_state_dict(torch.load(directory / "text_model.pt", map_location="cpu", weights_only=True)) |
| tokenizer = TinyTokenizer() |
| else: |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| tokenizer = AutoTokenizer.from_pretrained(directory, local_files_only=True) |
| if config["lora_rank"]: |
| from peft import PeftModel |
| base = build_text_model(config["backbone"], lora_rank=0, local_files_only=local_files_only, |
| backbone_dtype=config["backbone_dtype"]) |
| model = PeftModel.from_pretrained(base, directory / "model") |
| else: |
| model = AutoModelForCausalLM.from_pretrained(directory / "model", local_files_only=True) |
| return model.to(device), tokenizer |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--data", type=Path, required=True, help="audited training JSONL") |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--backbone", default="Qwen/Qwen2.5-0.5B-Instruct") |
| parser.add_argument("--steps", type=int, default=1000) |
| parser.add_argument("--batch-size", type=int, default=2) |
| parser.add_argument("--lr", type=float, default=1e-4) |
| parser.add_argument("--max-length", type=int, default=4096) |
| parser.add_argument("--hidden-size", type=int, default=64, help="tiny backend only") |
| parser.add_argument("--lora-rank", type=int, default=8) |
| parser.add_argument("--lora-alpha", type=int, default=16) |
| parser.add_argument("--backbone-dtype", choices=["float32", "float16", "bfloat16"], default="float32") |
| parser.add_argument("--device", default="cpu") |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--local-files-only", action="store_true") |
| parser.add_argument("--dry-run", action="store_true", help="one offline optimizer step; requires explicit tiny backend") |
| args = parser.parse_args(argv) |
| if args.steps < 1 or args.batch_size < 1 or not math.isfinite(args.lr) or args.lr <= 0: |
| parser.error("steps, batch size and learning rate must be positive and finite") |
| if args.dry_run and args.backbone != "tiny": |
| parser.error("offline dry-run requires --backbone tiny") |
| output = safe_output(args.output) |
| torch.manual_seed(args.seed) |
| metadata = run_metadata(args.data) |
| samples = _read_samples(args.data) |
| _verify_source_fingerprint(metadata) |
| tokenizer = load_tokenizer(args.backbone, local_files_only=args.local_files_only) |
| for row, sample in enumerate(samples): |
| try: |
| collate_text_samples([sample], tokenizer, max_length=args.max_length) |
| except ValueError as exc: |
| raise ValueError(f"text training row {row} failed preflight: {exc}") from exc |
| model = build_text_model(args.backbone, hidden_size=args.hidden_size, lora_rank=args.lora_rank, |
| lora_alpha=args.lora_alpha, local_files_only=args.local_files_only, |
| backbone_dtype=args.backbone_dtype).to(args.device).train() |
| optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr) |
| generator = torch.Generator().manual_seed(args.seed) |
| steps = 1 if args.dry_run else args.steps |
| log = [] |
| for step in range(steps): |
| indices = torch.randint(len(samples), (args.batch_size,), generator=generator).tolist() |
| batch = collate_text_samples([samples[i] for i in indices], tokenizer, max_length=args.max_length) |
| batch = {key: value.to(args.device) if isinstance(value, torch.Tensor) else value for key, value in batch.items()} |
| optimizer.zero_grad(set_to_none=True) |
| loss = text_loss(model, batch) |
| if not torch.isfinite(loss): |
| raise RuntimeError(f"non-finite assistant CE at step {step}") |
| loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True) |
| optimizer.step() |
| record = {"step": step + 1, "assistant_ce": float(loss.detach())} |
| log.append(record) |
| if step == 0 or (step + 1) % 10 == 0 or step + 1 == steps: |
| print(json.dumps(record), flush=True) |
| _verify_source_fingerprint(metadata) |
| safe_output(output, create=True) |
| tokenizer.save_pretrained(output) |
| if args.backbone == "tiny": |
| torch.save(model.state_dict(), output / "text_model.pt") |
| else: |
| model.save_pretrained(output / "model") |
| config = {key: value for key, value in vars(args).items() if key not in {"data", "output"}} |
| config.update(data=str(args.data), implementation="text_token_ce", offline_smoke=args.backbone == "tiny") |
| config["run_metadata"] = metadata |
| (output / "text_config.json").write_text(json.dumps(config, indent=2) + "\n") |
| (output / "text_training_log.json").write_text(json.dumps(log, indent=2) + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
