From 6a7196fc29ae5cd4ca4c325f2b154521477614d6 Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Mon, 24 Aug 2026 22:30:59 -0400 Subject: [PATCH 1/6] Harden training data and Gemma 4 mode handling --- docs/cli.md | 2 +- docs/generation.md | 2 +- docs/studio.md | 2 +- docs/training.md | 53 +++++++++- gemma4_example.py | 8 +- src/teich/anonymize.py | 90 ++++++++++++++-- src/teich/audit.py | 4 +- src/teich/cli.py | 3 + src/teich/config.py | 32 +++--- src/teich/formatter.py | 159 +++++++++++++++++++++++++++- src/teich/runner.py | 144 +++++++++++-------------- src/teich/studio/extraction.py | 29 ++++- src/teich/studio/generation.py | 14 ++- src/teich/studio/interactive.py | 33 +++++- src/teich/studio/server.py | 27 +++-- tests/test_audit.py | 11 ++ tests/test_cli.py | 14 +-- tests/test_config.py | 7 ++ tests/test_extract_anonymize_cli.py | 68 +++++++++--- tests/test_formatter.py | 120 +++++++++++++++++++++ tests/test_runner.py | 11 +- tests/test_studio.py | 69 ++++++++++++ tests/test_tokenizer_smoke.py | 63 ++++++++++- 23 files changed, 808 insertions(+), 157 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 1fd17b4..b853511 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -98,7 +98,7 @@ teich anonymize output --output output_anonymized teich anonymize data --in-place ``` -Anonymization replaces known credential formats, high-confidence secret assignments, personal email addresses, and home-directory usernames with deterministic dummy values while preserving embedded base64 media payloads. Reserved example-domain addresses, known public bot addresses, provider thinking signatures, placeholders, and common public IDs are preserved to avoid corrupting training data. Assistant-authored prose, reasoning, and code examples are also preserved from credential/email heuristics; local usernames in paths are still anonymized, and structured tool inputs continue through the full scanner. Reported totals are replacement occurrences rather than estimates of unique secrets. It is a best-effort pass; review data before publishing. +Anonymization replaces known credential formats, high-confidence secret assignments, personal email addresses, contextual PII, home-directory usernames, and embedded base64 media with deterministic dummy values. Reserved example-domain addresses, known public bot addresses, provider thinking signatures, placeholders, and common public IDs are preserved to avoid corrupting training data. User, assistant, reasoning, and structured tool content all pass through the same high-confidence privacy scanner. Reported totals are replacement occurrences rather than estimates of unique secrets. It is a best-effort pass; review data before publishing. ## Studio diff --git a/docs/generation.md b/docs/generation.md index 51c0d63..20bed03 100644 --- a/docs/generation.md +++ b/docs/generation.md @@ -81,7 +81,7 @@ teich extract codex --model gpt-5-codex --out codex-data `--model` filters by provider model metadata, not by arbitrary prompt text. This keeps traces that actually ran with matching model identifiers such as `claude-fable-5` and excludes traces that only mention the model name in conversation text. -After extraction, Teich automatically scrubs API keys, emails, and home-directory usernames while preserving embedded media payloads for conversation context. It then prints the replacement counts and asks whether to upload to Hugging Face. If you need a raw, unchanged local export, pass `--no-anon` or `--no-anonymize`: +After extraction, Teich automatically scrubs API keys, emails, contextual PII, home-directory usernames, and embedded media payloads. It then prints the replacement counts and asks whether to upload to Hugging Face. If you need a raw, unchanged local export, pass `--no-anon` or `--no-anonymize`: ```bash teich extract codex --sessions-dir /path/to/.codex --out raw-codex-data --no-anon diff --git a/docs/studio.md b/docs/studio.md index 160f88a..d2d1084 100644 --- a/docs/studio.md +++ b/docs/studio.md @@ -119,7 +119,7 @@ https://huggingface.co/datasets///embed/viewer That official embed works for datasets already available on the Hub. For unpublished local output, Studio uses Teich's local converter to approximate the parts of the viewer that matter before upload. The full Hugging Face viewer backend is hosted by Hugging Face and adds Parquet-backed row serving, search, filtering, SQL, and statistics after the dataset is uploaded and processed. -The Studio upload button regenerates the dataset card before publishing. The card stays intentionally short and points readers to the maintained training docs; large dataset-level tool snapshots are written to `tools.json` and uploaded alongside the JSONL files. +The Studio upload button regenerates the dataset card before publishing. The card stays intentionally short and points readers to the maintained training docs; dataset JSONL, generated metadata JSON, `README.md`, and `tools.json` are allowlisted for upload. Unrelated files in the output directory are not published. ## Requirements diff --git a/docs/training.md b/docs/training.md index 391a0d3..bf2ea9c 100644 --- a/docs/training.md +++ b/docs/training.md @@ -57,11 +57,54 @@ are: - `google/gemma-4-26B-A4B-it` - `google/gemma-4-31B-it` -Pass `chat_template_kwargs={"enable_thinking": True, "preserve_thinking": True}` -to `prepare_data()`. Do not replace `tokenizer.chat_template` unless you are -intentionally testing a maintained fork. Teich supervises the closing -`` token for completed Gemma responses while keeping system, user, and -tool-response context masked. +Gemma 4 does not infer this mode from the presence of reasoning labels alone. +The live template enables thinking by placing `<|think|>` at the beginning of +the system turn when `enable_thinking=True`; do not embed that token manually, +because enabling the template would then duplicate it. For thinking SFT, pass +`chat_template_kwargs={"enable_thinking": True, "preserve_thinking": True}`. +`preserve_thinking=True` is required for multi-turn rows with historical +reasoning; otherwise the upstream template silently omits reasoning from all +but the last assistant turn. + +For non-thinking SFT, pass `{"enable_thinking": False}` and use rows that have +neither reasoning fields nor a manually embedded `<|think|>` system trigger. +The live template will render `reasoning_content` even when thinking is off, +which creates an internally inconsistent example; Teich rejects that state. +The 26B-A4B and 31B generation prompts also insert an empty thought channel in +non-thinking mode while E4B does not. Teich detects this from the loaded live +template and mirrors the prefix in completed training turns so SFT and +inference contexts agree. + +Do not replace `tokenizer.chat_template` unless you are intentionally testing a +maintained fork. Teich supervises the closing `` token for completed +Gemma responses while keeping system, user, tool-response, and generated prompt +prefix context masked. + +Thinking and non-thinking examples can be mixed safely by separating them into +sources with source-level template kwargs: + +```python +train_dataset = prepare_data( + { + "thinking": { + "source": "username/gemma-thinking-traces", + "percentage": 60, + "chat_template_kwargs": { + "enable_thinking": True, + "preserve_thinking": True, + }, + }, + "direct": { + "source": "username/gemma-direct-traces", + "percentage": 40, + "chat_template_kwargs": {"enable_thinking": False}, + }, + }, + tokenizer, + tokenize=True, + strict=True, +) +``` `gemma4_example.py` uses the live remote template by default. Set `CHAT_TEMPLATE_PATH` only to opt into a local custom template, and set diff --git a/gemma4_example.py b/gemma4_example.py index 5667aa3..427663d 100644 --- a/gemma4_example.py +++ b/gemma4_example.py @@ -17,6 +17,7 @@ HUB_REPO_ID = os.environ.get("HUB_REPO_ID") or "" HF_TOKEN = os.environ.get("HF_TOKEN", "") CHAT_TEMPLATE_PATH = os.environ.get("CHAT_TEMPLATE_PATH") +ENABLE_THINKING = os.environ.get("GEMMA4_ENABLE_THINKING", "1").strip().lower() not in {"0", "false", "no"} model, tokenizer = FastModel.from_pretrained( model_name=MODEL_NAME, @@ -66,7 +67,12 @@ tokenizer, split="train", hf_token=HF_TOKEN, - chat_template_kwargs={"enable_thinking": True, "preserve_thinking": True}, + # Non-thinking datasets must not contain reasoning fields or a manual + # <|think|> system trigger; Teich rejects those inconsistent rows. + chat_template_kwargs={ + "enable_thinking": ENABLE_THINKING, + "preserve_thinking": ENABLE_THINKING, + }, max_length=MAX_SEQ_LEN, oversized_policy="trim_followups", tokenize=True, diff --git a/src/teich/anonymize.py b/src/teich/anonymize.py index f06918a..217fe3b 100644 --- a/src/teich/anonymize.py +++ b/src/teich/anonymize.py @@ -5,6 +5,7 @@ from concurrent.futures import ProcessPoolExecutor, as_completed from dataclasses import dataclass, field from pathlib import Path +import base64 import hashlib import multiprocessing import os @@ -286,6 +287,25 @@ class TraceAnonymizer: r"(?:(?\b(?:phone|telephone|tel|call(?:\s+me)?\s+at)\s*[:=]?\s*)" + r"(?P(?:\+?1[\s.-]?)?(?:(?:\(\d{3}\)|\d{3})[\s.-]?)\d{3}[\s.-]?\d{4})" + ) + _ip_pattern = re.compile( + r"(?i)(?P\b(?:ip|ip\s+address|host\s+address)\s*[:=]?\s*)" + r"(?P(?:\d{1,3}\.){3}\d{1,3})" + ) + _name_pattern = re.compile( + r"(?i)(?P\b(?:my\s+name\s+is|full\s+name\s*[:=]|name\s*[:=])\s*)" + r"(?P[A-Z][A-Za-z'-]+(?:\s+[A-Z][A-Za-z'-]+){1,3})" + ) + _address_pattern = re.compile( + r"(?i)(?P\b(?:address\s*[:=]|i\s+live\s+at)\s*)" + r"(?P\d{1,6}\s+[A-Za-z0-9][A-Za-z0-9 .'-]{1,80}" + r"(?:street|st|road|rd|avenue|ave|boulevard|blvd|lane|ln|drive|dr|court|ct|way)\b" + r"(?:,\s*[A-Za-z .'-]{2,40})?(?:,\s*[A-Z]{2}\s+\d{5}(?:-\d{4})?)?)" + ) # Match an assignment name from the start of the complete identifier, then # decide whether it is sensitive in _replace_generic_secret. The explicit # identifier-character lookbehind is load-bearing: unlike a word boundary, @@ -642,8 +662,24 @@ class TraceAnonymizer: "dsn", "connection", ) + _structured_pii_keys = { + "address", + "first_name", + "full_name", + "home_address", + "ip_address", + "last_name", + "mobile", + "phone", + "phone_number", + "social_security_number", + "ssn", + "street_address", + "telephone", + } + def __init__(self) -> None: - self.counts = {"email": 0, "username": 0, "api_key": 0} + self.counts = {"email": 0, "username": 0, "api_key": 0, "pii": 0, "media": 0} self._email_map: dict[str, str] = {} self._username_map: dict[str, str] = {} self._api_key_map: dict[str, str] = {} @@ -660,26 +696,26 @@ def anonymize_value(self, value: Any) -> Any: return value def _anonymize_mapping(self, value: dict[Any, Any]) -> dict[Any, Any]: - should_preserve_base64_data = self._looks_like_base64_media_source(value) + should_redact_base64_data = self._looks_like_base64_media_source(value) role = value.get("role") is_assistant_message = isinstance(role, str) and role.lower() == "assistant" redacted: dict[Any, Any] = {} for key, item in value.items(): redacted_key = self.anonymize_value(key) if ( - should_preserve_base64_data + should_redact_base64_data and key == "data" and isinstance(item, str) and self._looks_like_base64_blob(item) ): - redacted[redacted_key] = item + redacted[redacted_key] = self._redact_base64_media(item) elif is_assistant_message and key in {"content", "reasoning_content", "thinking"}: - # Assistant-authored prose and code frequently demonstrate - # credentials, JWTs, connection strings, and email addresses. - # Treating those examples as private input corrupts otherwise - # valid training targets. Identity-bearing local paths still - # get scrubbed; structured tool inputs remain on the full path. + # Assistant responses can echo credentials and personal data + # supplied by users or tools. Apply the same high-confidence + # privacy rules here instead of assuming the content is synthetic. redacted[redacted_key] = self._anonymize_assistant_generated_value(item) + elif isinstance(key, str) and isinstance(item, str) and self._is_structured_pii_key(key): + redacted[redacted_key] = self._redact_structured_pii(item) elif isinstance(key, str) and self._should_redact_mapping_value(key, item): redacted[redacted_key] = self._redact_mapping_value(key, item) else: @@ -688,7 +724,7 @@ def _anonymize_mapping(self, value: dict[Any, Any]) -> dict[Any, Any]: def _anonymize_assistant_generated_value(self, value: Any) -> Any: if isinstance(value, str): - return self._anonymize_identity_text(value) + return self.anonymize_text(value) if isinstance(value, list): return [self._anonymize_assistant_content_block(item) for item in value] if isinstance(value, dict): @@ -727,6 +763,15 @@ def _redact_mapping_value(self, key: str, item: Any) -> str: value = str(item) return self._assignment_secret_replacement(key, value) + @classmethod + def _is_structured_pii_key(cls, key: str) -> bool: + normalized = re.sub(r"[^a-z0-9]+", "_", key.lower()).strip("_") + return normalized in cls._structured_pii_keys + + def _redact_structured_pii(self, value: str) -> str: + self.counts["pii"] += 1 + return self._pii_replacement(value) + @staticmethod def _looks_like_base64_media_source(value: dict[Any, Any]) -> bool: source_type = value.get("type") @@ -747,14 +792,39 @@ def _looks_like_base64_blob(value: str) -> bool: return False return re.fullmatch(r"[A-Za-z0-9+/=\s]+", value) is not None + def _redact_base64_media(self, value: str) -> str: + """Replace opaque media with a deterministic, valid base64 marker.""" + marker = f"[redacted media {self._dummy_sequence(value, 16)}]".encode("utf-8") + self.counts["media"] += 1 + return base64.b64encode(marker).decode("ascii") + def anonymize_text(self, text: str) -> str: lowered = text.lower() if "@" in text: text = self._replace_emails(text) text = self._anonymize_identity_text(text, lowered=lowered) text = self._replace_api_keys(text) + text = self._replace_high_confidence_pii(text) + return text + + def _replace_high_confidence_pii(self, text: str) -> str: + def replace_value(match: re.Match[str]) -> str: + value = match.group("value") + self.counts["pii"] += 1 + return match.group("prefix") + self._pii_replacement(value) + + def replace_ssn(match: re.Match[str]) -> str: + self.counts["pii"] += 1 + return self._pii_replacement(match.group(0)) + + text = self._ssn_pattern.sub(replace_ssn, text) + for pattern in (self._phone_pattern, self._ip_pattern, self._name_pattern, self._address_pattern): + text = pattern.sub(replace_value, text) return text + def _pii_replacement(self, value: str) -> str: + return "redacted_pii_" + self._dummy_sequence(value, 16) + def _anonymize_identity_text(self, text: str, *, lowered: str | None = None) -> str: lowered = text.lower() if lowered is None else lowered if "home" in lowered or "users" in lowered: diff --git a/src/teich/audit.py b/src/teich/audit.py index e253b12..6b8d3da 100644 --- a/src/teich/audit.py +++ b/src/teich/audit.py @@ -96,7 +96,7 @@ def _audit_training_row(row: dict[str, Any], tokenizer: Any, row_index: int) -> return errors, warnings, sample -def audit_sft_dataset(dataset: Dataset, tokenizer: Any, *, sample_size: int = 8) -> SFTAuditReport: +def audit_sft_dataset(dataset: Dataset, tokenizer: Any, *, sample_size: int | None = None) -> SFTAuditReport: if not isinstance(dataset, Dataset): return SFTAuditReport(ok=False, errors=["dataset must be a datasets.Dataset instance"]) errors: list[str] = [] @@ -111,7 +111,7 @@ def audit_sft_dataset(dataset: Dataset, tokenizer: Any, *, sample_size: int = 8) if dataset.num_rows == 0: return SFTAuditReport(ok=False, errors=["dataset contains no rows"]) - limit = min(max(sample_size, 0), dataset.num_rows) + limit = dataset.num_rows if sample_size is None else min(max(sample_size, 0), dataset.num_rows) if limit == 0: warnings.append("sample_size is 0; no rows audited") diff --git a/src/teich/cli.py b/src/teich/cli.py index 42f2041..0acbc7d 100644 --- a/src/teich/cli.py +++ b/src/teich/cli.py @@ -192,6 +192,7 @@ def _help_group(name: str, extra_help: str) -> type[ExtraHelpGroup]: NON_DATA_TRACE_DIR_NAMES = {"partials", "failures"} UPLOAD_IGNORE_PATTERNS = ["partials/**", "failures/**"] UPLOAD_METADATA_PATTERNS = ["README.md", "tools.json"] +UPLOAD_DATA_PATTERNS = ["*.jsonl", "**/*.jsonl", "*.metadata.json", "**/*.metadata.json"] def _upload_ignore_patterns(cfg: Config) -> list[str]: @@ -326,6 +327,7 @@ def _upload_dataset_folder( folder_path=str(folder_path), repo_type="dataset", private=private, + allow_patterns=UPLOAD_DATA_PATTERNS, ignore_patterns=list(dict.fromkeys([*ignore_patterns, *UPLOAD_METADATA_PATTERNS])), ) else: @@ -334,6 +336,7 @@ def _upload_dataset_folder( repo_id=repo_id, repo_type="dataset", commit_message="Upload teich dataset output", + allow_patterns=[*UPLOAD_DATA_PATTERNS, *UPLOAD_METADATA_PATTERNS], ignore_patterns=ignore_patterns, ) return str(repo_url) diff --git a/src/teich/config.py b/src/teich/config.py index 3c9c8d2..7cceb71 100644 --- a/src/teich/config.py +++ b/src/teich/config.py @@ -11,7 +11,7 @@ from typing import Any, cast import yaml -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator GITHUB_REPO_ID_PATTERN = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") @@ -63,7 +63,13 @@ def _raise_csv_field_limit() -> None: limit //= 10 -class MCPConfig(BaseModel): +class StrictConfigModel(BaseModel): + """Base for user-authored configuration that rejects misspelled keys.""" + + model_config = ConfigDict(extra="forbid") + + +class MCPConfig(StrictConfigModel): """MCP server configuration.""" name: str command: str | None = None @@ -95,7 +101,7 @@ def validate_transport(self) -> MCPConfig: return self -class APIConfig(BaseModel): +class APIConfig(StrictConfigModel): """API configuration for OpenAI-compatible endpoints.""" provider: str = "openai" # openai, openrouter, azure, etc. base_url: str | None = None # e.g., https://openrouter.ai/api/v1 @@ -103,7 +109,7 @@ class APIConfig(BaseModel): wire_api: str = "responses" -class LangfuseConfig(BaseModel): +class LangfuseConfig(StrictConfigModel): """Langfuse tracing credentials, set under ``agent.langfuse``. When enabled, Teich wires each agent's Langfuse integration (the Codex @@ -135,7 +141,7 @@ def require_credentials_when_enabled(self) -> LangfuseConfig: return self -class CodexAuthConfig(BaseModel): +class CodexAuthConfig(StrictConfigModel): """Codex ChatGPT-subscription auth handling. When ``use_host_auth`` is enabled, Teich seeds an ``auth.json`` snapshot @@ -155,7 +161,7 @@ class CodexAuthConfig(BaseModel): broker_port: int = Field(default=0, ge=0, le=65535) -class ClaudeConfig(BaseModel): +class ClaudeConfig(StrictConfigModel): """Claude Code-specific settings, set under ``agent.claude``. Subscription auth (Pro/Max): when a long-lived OAuth token is available — @@ -190,7 +196,7 @@ class ClaudeConfig(BaseModel): max_thinking_tokens: int | None = Field(default=None, ge=0) -class HarnessContextCaptureConfig(BaseModel): +class HarnessContextCaptureConfig(StrictConfigModel): """Safe simulated capture of client-visible harness instructions and tools. When enabled, Teich points the configured harness at a local fake provider @@ -204,7 +210,7 @@ class HarnessContextCaptureConfig(BaseModel): timeout_seconds: int = Field(default=45, gt=0, le=300) -class AgentConfig(BaseModel): +class AgentConfig(StrictConfigModel): """Agent runtime selection.""" provider: str = "codex" # Langfuse tracing, applied to every agent that supports it (Codex, Claude). @@ -213,7 +219,7 @@ class AgentConfig(BaseModel): claude: ClaudeConfig = Field(default_factory=ClaudeConfig) -class ModelConfig(BaseModel): +class ModelConfig(StrictConfigModel): """Model configuration.""" model: str = "codex-mini-latest" approval_policy: str = "never" @@ -240,7 +246,7 @@ def normalize_legacy_approval_mode(self) -> ModelConfig: return self -class OutputConfig(BaseModel): +class OutputConfig(StrictConfigModel): """Output configuration.""" traces_dir: Path = Field(default=Path("./output")) sandbox_dir: Path = Field(default=Path("./sandbox")) @@ -248,7 +254,7 @@ class OutputConfig(BaseModel): pretty_name: str = "Agentic Training Traces" -class PublishConfig(BaseModel): +class PublishConfig(StrictConfigModel): """Publishing configuration.""" repo_id: str | None = None hf_token: str | None = None @@ -267,7 +273,7 @@ def validate_repo_id(cls, value: str | None) -> str | None: return normalized -class PromptInput(BaseModel): +class PromptInput(StrictConfigModel): """Structured prompt input row.""" image: str | None = None github_repo: str | None = None @@ -337,7 +343,7 @@ def turn_prompts(self) -> list[str]: return [self.prompt, *self.follow_up_prompts] -class Config(BaseModel): +class Config(StrictConfigModel): """Main configuration.""" agent: AgentConfig = Field(default_factory=AgentConfig) model: ModelConfig = Field(default_factory=ModelConfig) diff --git a/src/teich/formatter.py b/src/teich/formatter.py index 3c80386..f8c70df 100644 --- a/src/teich/formatter.py +++ b/src/teich/formatter.py @@ -5,6 +5,7 @@ from dataclasses import dataclass, field import json import re +import weakref from collections.abc import Mapping, Sequence from typing import Any @@ -18,6 +19,8 @@ _GEMMA_ASSISTANT_TURN_PREFIX = "<|turn>model\n" _GEMMA_TURN_END = "" _GEMMA_THOUGHT_PREFIX = "<|channel>thought\n" +_GEMMA_THOUGHT_END = "" +_GEMMA_THINK_TRIGGER = "<|think|>" _GEMMA_TOOL_RESPONSE_START = "<|tool_response>" _GEMMA_TOOL_RESPONSE_END = "" _TOOL_RESPONSE_DELIMITERS = ( @@ -72,6 +75,7 @@ } _OVERSIZED_POLICIES = {"drop", "trim_followups", "error"} _OVERSIZED_POLICY_KEEP = "keep" +_GEMMA4_PREFIX_CACHE: weakref.WeakKeyDictionary[Any, dict[str, str]] = weakref.WeakKeyDictionary() @dataclass(slots=True) @@ -208,6 +212,156 @@ def _validate_chat_template_kwargs(chat_template_kwargs: dict[str, Any] | None) return kwargs +def _renderer_name(renderer: Any) -> str: + for candidate in (renderer, getattr(renderer, "tokenizer", None)): + name = getattr(candidate, "name_or_path", None) + if isinstance(name, str) and name: + return name.lower() + return "" + + +def _is_gemma4_renderer(renderer: Any) -> bool: + if "gemma-4" in _renderer_name(renderer): + return True + template = getattr(renderer, "chat_template", None) + return ( + isinstance(template, str) + and "<|turn>model" in template + and _GEMMA_THINK_TRIGGER in template + ) + + +def _message_reasoning(message: dict[str, Any]) -> str: + for key in ("reasoning_content", "thinking", "reasoning"): + value = message.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + +def _message_text_contains(message: dict[str, Any], needle: str) -> bool: + content = message.get("content") + if isinstance(content, str): + return needle in content + if isinstance(content, list): + return any( + isinstance(part, dict) + and isinstance(part.get("text"), str) + and needle in part["text"] + for part in content + ) + return False + + +def _validate_gemma4_thinking_contract( + renderer: Any, + messages: list[dict[str, Any]], + chat_template_kwargs: dict[str, Any], +) -> None: + if not _is_gemma4_renderer(renderer): + return + thinking_enabled = chat_template_kwargs.get("enable_thinking", False) is True + reasoning_indexes = [ + index + for index, message in enumerate(messages) + if message.get("role") in {"assistant", "model"} and _message_reasoning(message) + ] + manual_trigger = any( + message.get("role") in {"system", "developer"} + and _message_text_contains(message, _GEMMA_THINK_TRIGGER) + for message in messages + ) + if manual_trigger: + raise ValueError( + "Gemma 4 system/developer content already contains <|think|>. Do not embed the trigger " + "manually: use chat_template_kwargs={'enable_thinking': True, " + "'preserve_thinking': True} so the live template inserts it exactly once." + ) + if not thinking_enabled and reasoning_indexes: + raise ValueError( + "Gemma 4 thinking is disabled, but this row contains reasoning. Use " + "chat_template_kwargs={'enable_thinking': True, " + "'preserve_thinking': True} for thinking data, or remove reasoning for non-thinking data." + ) + if thinking_enabled and chat_template_kwargs.get("preserve_thinking", False) is not True: + last_assistant = max( + ( + index + for index, message in enumerate(messages) + if message.get("role") in {"assistant", "model"} + ), + default=-1, + ) + if any(index != last_assistant for index in reasoning_indexes): + raise ValueError( + "Gemma 4 preserve_thinking=False drops reasoning from earlier assistant turns. " + "Set preserve_thinking=True for multi-turn thinking SFT." + ) + + +def _gemma4_nonthinking_generation_suffix( + renderer: Any, + chat_template_kwargs: dict[str, Any], +) -> str: + if not _is_gemma4_renderer(renderer) or chat_template_kwargs.get("enable_thinking", False) is True: + return "" + try: + serialized_kwargs = json.dumps(chat_template_kwargs, sort_keys=True, default=repr) + except TypeError: + serialized_kwargs = repr(chat_template_kwargs) + try: + renderer_cache = _GEMMA4_PREFIX_CACHE.setdefault(renderer, {}) + except TypeError: + renderer_cache = {} + if serialized_kwargs in renderer_cache: + return renderer_cache[serialized_kwargs] + probe = [{"role": "user", "content": "__TEICH_GEMMA4_PROBE__"}] + base_kwargs = {"tokenize": False, "add_generation_prompt": False, **chat_template_kwargs} + prompt_kwargs = {"tokenize": False, "add_generation_prompt": True, **chat_template_kwargs} + try: + base = _apply_chat_template_with_gemma_fallback(renderer, probe, base_kwargs) + prompt = _apply_chat_template_with_gemma_fallback(renderer, probe, prompt_kwargs) + except Exception: + renderer_cache[serialized_kwargs] = "" + return "" + if not isinstance(base, str) or not isinstance(prompt, str) or not prompt.startswith(base): + renderer_cache[serialized_kwargs] = "" + return "" + generation_prefix = prompt[len(base):] + if not generation_prefix.startswith(_GEMMA_ASSISTANT_TURN_PREFIX): + renderer_cache[serialized_kwargs] = "" + return "" + suffix = generation_prefix[len(_GEMMA_ASSISTANT_TURN_PREFIX):] + if suffix == _GEMMA_THOUGHT_PREFIX + _GEMMA_THOUGHT_END: + renderer_cache[serialized_kwargs] = suffix + return suffix + renderer_cache[serialized_kwargs] = "" + return "" + + +def _align_gemma4_nonthinking_training_text( + renderer: Any, + text: str, + chat_template_kwargs: dict[str, Any], +) -> str: + suffix = _gemma4_nonthinking_generation_suffix(renderer, chat_template_kwargs) + if not suffix: + return text + parts: list[str] = [] + cursor = 0 + for match in _GEMMA_TURN_START_PATTERN.finditer(text): + if match.group(1) != "model": + continue + parts.append(text[cursor:match.end()]) + if not text.startswith(_GEMMA_THOUGHT_PREFIX, match.end()): + parts.append(suffix) + cursor = match.end() + if not parts: + return text + parts.append(text[cursor:]) + return "".join(parts) + + def _as_text_content_parts(content: Any) -> Any: if isinstance(content, str): return [{"type": "text", "text": content}] @@ -284,7 +438,7 @@ def _render_chat( rendered = _apply_chat_template_with_gemma_fallback(renderer, messages, render_kwargs) if not isinstance(rendered, str): raise TypeError("tokenizer.apply_chat_template(..., tokenize=False) must return a string") - return rendered + return _align_gemma4_nonthinking_training_text(renderer, rendered, chat_template_kwargs) def _normalize_tool_call_arguments_for_template(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -1980,6 +2134,7 @@ def _render_training_row( assistant_prompt_prefix_cache: dict[str, tuple[str, ...]], strict: bool, ) -> _RenderedRow | None: + _validate_gemma4_thinking_contract(renderer, messages, template_kwargs) if teich_masking: text, supervised_spans = _supervised_text_and_spans( renderer, @@ -2550,7 +2705,7 @@ def mask_data( train_on_tool_responses: bool = False, max_supervised_tokens: int | None = None, audit: bool = True, - audit_sample_size: int = 8, + audit_sample_size: int | None = None, verbose: bool = True, ) -> Any: from .audit import audit_sft_dataset diff --git a/src/teich/runner.py b/src/teich/runner.py index 2b0a3a0..5c5b2ed 100644 --- a/src/teich/runner.py +++ b/src/teich/runner.py @@ -626,14 +626,7 @@ def _prompt_completion_key(prompt_input: PromptInput | str) -> str: def _prompt_resume_key(prompt_input: PromptInput | str) -> str: - if isinstance(prompt_input, str): - prompt_parts = [_prompt_text_completion_key(prompt_input)] - else: - prompt_parts = [ - _prompt_text_completion_key(prompt) - for prompt in prompt_input.turn_prompts() - ] - return "\n\n--- follow-up ---\n\n".join(prompt_parts) + return _prompt_completion_key(prompt_input) def _agent_turn_prompts(prompt: str, prompt_input: PromptInput | None) -> list[str]: @@ -883,6 +876,36 @@ def _system_from_training_example(example: dict[str, Any]) -> str | None: return None +def _system_from_teich_trace_events(events: list[dict[str, Any]]) -> str | None: + for event in reversed(events): + if event.get("type") != "custom" or event.get("customType") != PI_SYSTEM_PROMPT_CUSTOM_TYPE: + continue + data = event.get("data") + system = data.get("systemPrompt") if isinstance(data, dict) else None + if isinstance(system, str) and system.strip(): + return system.strip() + return None + + +def _trace_includes_system_prompt(events: list[dict[str, Any]], system_prompt: str) -> bool: + expected = system_prompt.strip() + if _system_from_teich_trace_events(events) == expected: + return True + for event in events: + candidates = [event] + candidates.extend( + value + for key in ("message", "payload") + if isinstance((value := event.get(key)), dict) + ) + for candidate in candidates: + if candidate.get("role") not in {"system", "developer"}: + continue + if _message_text(candidate.get("content")) == expected: + return True + return False + + def completed_prompt_keys_from_outputs(traces_dir: Path, excluded_dirs: list[Path] | None = None) -> set[str]: if not traces_dir.exists(): return set() @@ -897,7 +920,12 @@ def completed_prompt_keys_from_outputs(traces_dir: Path, excluded_dirs: list[Pat if structured_rows is not None: examples = structured_rows else: - examples = [convert_trace_to_training_example(path).to_dict()] + events = _read_jsonl_dict_events(path) or [] + example = convert_trace_to_training_example(path).to_dict() + recorded_system = _system_from_teich_trace_events(events) + if recorded_system and not _system_from_training_example(example): + example["system"] = recorded_system + examples = [example] except (OSError, json.JSONDecodeError, ValueError): continue for example in examples: @@ -1030,8 +1058,32 @@ def _append_harness_context_metadata(self, trace_path: Path) -> None: return self._append_jsonl_events(trace_path, [capture.to_trace_event()]) + def _append_prompt_system_metadata(self, trace_path: Path, prompt_input: PromptInput | None) -> None: + if prompt_input is None or not isinstance(prompt_input.system, str) or not trace_path.exists(): + return + system_prompt = prompt_input.system.strip() + if not system_prompt: + return + events = _read_jsonl_dict_events(trace_path) or [] + if _trace_includes_system_prompt(events, system_prompt): + return + self._append_jsonl_events( + trace_path, + [ + { + "type": "custom", + "id": f"teich-system-{uuid.uuid4().hex[:8]}", + "parentId": None, + "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"), + "customType": PI_SYSTEM_PROMPT_CUSTOM_TYPE, + "data": {"systemPrompt": system_prompt, "source": "teich"}, + } + ], + ) + def _finalize_trace_export(self, trace_path: Path, prompt_input: PromptInput | None = None) -> None: self._append_harness_context_metadata(trace_path) + self._append_prompt_system_metadata(trace_path, prompt_input) def capture_harness_context(self) -> HarnessContextCapture: raise RuntimeError( @@ -6168,78 +6220,6 @@ def _copy_normalized_session_file(self, source_path: Path, destination: Path) -> destination.parent.mkdir(parents=True, exist_ok=True) shutil.copyfile(source_path, destination) - @staticmethod - def _prompt_input_system_prompt(prompt_input: PromptInput | None) -> str | None: - if prompt_input is None or not isinstance(prompt_input.system, str): - return None - system_prompt = prompt_input.system.strip() - return system_prompt or None - - @classmethod - def _pi_trace_includes_system_prompt(cls, trace_path: Path, system_prompt: str) -> bool: - expected = system_prompt.strip() - if not expected: - return True - try: - with trace_path.open("r", encoding="utf-8") as source: - for raw_line in source: - line = raw_line.strip() - if not line: - continue - event = json.loads(line) - if not isinstance(event, dict): - continue - if event.get("type") == "custom" and event.get("customType") == PI_SYSTEM_PROMPT_CUSTOM_TYPE: - data = event.get("data") - recorded = data.get("systemPrompt") if isinstance(data, dict) else None - if isinstance(recorded, str) and recorded.strip() == expected: - return True - continue - if event.get("type") != "message": - continue - payload = event.get("message") - if not isinstance(payload, dict) or payload.get("role") not in {"system", "developer"}: - continue - if cls._pi_first_text(payload.get("content")).strip() == expected: - return True - except (OSError, json.JSONDecodeError): - return False - return False - - @staticmethod - def _pi_first_text(content: Any) -> str: - if isinstance(content, str): - return content - if not isinstance(content, list): - return "" - for block in content: - if not isinstance(block, dict): - continue - if block.get("type") != "text": - continue - text = block.get("text") - if isinstance(text, str): - return text - return "" - - def _append_pi_system_prompt_metadata(self, trace_path: Path, prompt_input: PromptInput | None) -> None: - system_prompt = self._prompt_input_system_prompt(prompt_input) - if not system_prompt or self._pi_trace_includes_system_prompt(trace_path, system_prompt): - return - event = { - "type": "custom", - "id": f"teich-system-{uuid.uuid4().hex[:8]}", - "parentId": None, - "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"), - "customType": PI_SYSTEM_PROMPT_CUSTOM_TYPE, - "data": { - "systemPrompt": system_prompt, - "source": "teich", - }, - } - with trace_path.open("a", encoding="utf-8") as destination: - destination.write(json.dumps(event, separators=(",", ":")) + "\n") - @staticmethod def _pi_trace_includes_available_tools(trace_path: Path) -> bool: try: @@ -6280,7 +6260,7 @@ def _append_pi_available_tools_metadata(self, trace_path: Path) -> None: destination.write(json.dumps(event, separators=(",", ":"), ensure_ascii=False) + "\n") def _finalize_trace_export(self, trace_path: Path, prompt_input: PromptInput | None = None) -> None: - self._append_pi_system_prompt_metadata(trace_path, prompt_input) + super()._finalize_trace_export(trace_path, prompt_input) self._append_pi_available_tools_metadata(trace_path) def _latest_session_source_file(self, session_id: str, session_dir: Path, started_at: datetime) -> Path: diff --git a/src/teich/studio/extraction.py b/src/teich/studio/extraction.py index cd57f3a..dea99be 100644 --- a/src/teich/studio/extraction.py +++ b/src/teich/studio/extraction.py @@ -7,6 +7,7 @@ import uuid from datetime import datetime, timezone from pathlib import Path +from collections.abc import Callable from typing import Any from ..config import Config @@ -14,6 +15,8 @@ from ..trace_readme import extraction_readme_tags, write_traces_readme from .interactive import EventLog +JOB_HISTORY_LIMIT = 20 + def _extract_dataset_config( provider: ExtractProvider, @@ -63,9 +66,16 @@ def __init__( self.detected_sources: list[str] = [] self.anonymize_totals: dict[str, int] | None = None self._lock = threading.RLock() + self._thread: threading.Thread | None = None def start(self) -> None: - threading.Thread(target=self._run, name=f"studio-extract-{self.id[:8]}", daemon=True).start() + self._thread = threading.Thread(target=self._run, name=f"studio-extract-{self.id[:8]}", daemon=True) + self._thread.start() + + def join(self) -> None: + thread = self._thread + if thread is not None and thread is not threading.current_thread(): + thread.join() def _emit_status(self, status: str, message: str | None = None) -> None: with self._lock: @@ -234,10 +244,14 @@ def start( source_paths: list[Path] | None = None, model_filter: str | None = None, skip_anonymize: bool = False, + before_start: Callable[[], None] | None = None, ) -> ExtractionJob: with self._lock: if self._current is not None and self._current.status in {"starting", "running"}: raise RuntimeError("An extraction run is already in progress") + if before_start is not None: + before_start() + self._prune_locked() job = ExtractionJob( provider, output_dir=output_dir, @@ -250,6 +264,14 @@ def start( job.start() return job + def _prune_locked(self) -> None: + overflow = len(self._jobs) - JOB_HISTORY_LIMIT + 1 + if overflow <= 0: + return + completed = [job_id for job_id, job in self._jobs.items() if job.status not in {"starting", "running"}] + for job_id in completed[:overflow]: + self._jobs.pop(job_id, None) + def current(self) -> ExtractionJob | None: with self._lock: return self._current @@ -262,7 +284,10 @@ def get(self, job_id: str) -> ExtractionJob: return job def shutdown(self) -> None: - pass + with self._lock: + job = self._current + if job is not None and job.status in {"starting", "running"}: + job.join() def _jsonl_row_count(paths: list[Path]) -> int: diff --git a/src/teich/studio/generation.py b/src/teich/studio/generation.py index f1b426d..9d2ea65 100644 --- a/src/teich/studio/generation.py +++ b/src/teich/studio/generation.py @@ -32,6 +32,7 @@ "hermes_agent": HermesRunner, "chat": ChatRunner, } +JOB_HISTORY_LIMIT = 20 class GenerationStopped(RuntimeError): @@ -66,9 +67,11 @@ def __init__(self, config: Config, *, resume: bool = False): self._lock = threading.RLock() self._runner: Any = None self._stop_requested = False + self._thread: threading.Thread | None = None def start(self) -> None: - threading.Thread(target=self._run, name=f"studio-generate-{self.id[:8]}", daemon=True).start() + self._thread = threading.Thread(target=self._run, name=f"studio-generate-{self.id[:8]}", daemon=True) + self._thread.start() def _emit_status(self, status: str, message: str | None = None) -> None: self.status = status @@ -230,12 +233,21 @@ def start(self, config: Config, *, resume: bool = False) -> GenerationJob: with self._lock: if self._current is not None and self._current.status in {"starting", "running"}: raise RuntimeError("A generation run is already in progress") + self._prune_locked() job = GenerationJob(config, resume=resume) self._jobs[job.id] = job self._current = job job.start() return job + def _prune_locked(self) -> None: + overflow = len(self._jobs) - JOB_HISTORY_LIMIT + 1 + if overflow <= 0: + return + completed = [job_id for job_id, job in self._jobs.items() if job.status not in {"starting", "running"}] + for job_id in completed[:overflow]: + self._jobs.pop(job_id, None) + def current(self) -> GenerationJob | None: with self._lock: return self._current diff --git a/src/teich/studio/interactive.py b/src/teich/studio/interactive.py index b587d9f..37710aa 100644 --- a/src/teich/studio/interactive.py +++ b/src/teich/studio/interactive.py @@ -55,20 +55,30 @@ } SCROLLBACK_LIMIT = 2 * 1024 * 1024 # bytes of terminal history kept for reconnects +EVENT_LOG_LIMIT = 10_000 +SESSION_HISTORY_LIMIT = 50 class EventLog: """Append-only event list with a condition for SSE long-polling.""" - def __init__(self) -> None: + def __init__(self, *, max_events: int = EVENT_LOG_LIMIT) -> None: self._events: list[dict[str, Any]] = [] self._condition = threading.Condition() + self._first_seq = 0 + self._next_seq = 0 + self._max_events = max(1, max_events) self.closed = False def append(self, event: dict[str, Any]) -> None: with self._condition: - event = {**event, "seq": len(self._events), "ts": time.time()} + event = {**event, "seq": self._next_seq, "ts": time.time()} + self._next_seq += 1 self._events.append(event) + overflow = len(self._events) - self._max_events + if overflow > 0: + del self._events[:overflow] + self._first_seq += overflow self._condition.notify_all() def close(self) -> None: @@ -79,9 +89,11 @@ def close(self) -> None: def wait_for(self, start_index: int, timeout: float = 15.0) -> list[dict[str, Any]]: """Return events at/after start_index, blocking up to timeout if none yet.""" with self._condition: - if len(self._events) <= start_index and not self.closed: + start_index = max(start_index, self._first_seq) + if self._next_seq <= start_index and not self.closed: self._condition.wait(timeout=timeout) - return self._events[start_index:] + start_index = max(start_index, self._first_seq) + return self._events[start_index - self._first_seq:] def snapshot(self) -> list[dict[str, Any]]: with self._condition: @@ -653,6 +665,7 @@ def __init__(self) -> None: def create(self, config: Config, *, github_repo: str | None = None, system: str | None = None) -> InteractiveSession: session = InteractiveSession(config, github_repo=github_repo, system=system) with self._lock: + self._prune_locked() self._sessions[session.id] = session session.start_async() return session @@ -673,6 +686,18 @@ def remove(self, session_id: str) -> None: with self._lock: self._sessions.pop(session_id, None) + def _prune_locked(self) -> None: + overflow = len(self._sessions) - SESSION_HISTORY_LIMIT + 1 + if overflow <= 0: + return + completed = [ + session_id + for session_id, session in self._sessions.items() + if session.status in {"finished", "saved", "discarded", "error"} + ] + for session_id in completed[:overflow]: + self._sessions.pop(session_id, None) + def shutdown(self) -> None: with self._lock: sessions = list(self._sessions.values()) diff --git a/src/teich/studio/server.py b/src/teich/studio/server.py index dc321f4..981df56 100644 --- a/src/teich/studio/server.py +++ b/src/teich/studio/server.py @@ -115,9 +115,11 @@ class DatasetUploadRequest(BaseModel): _docker_cache: dict[str, Any] = {"checked_at": 0.0, "available": False, "detail": None} TERMINAL_READY_STATUSES = {"ready", "live", "exited", "error"} TERMINAL_STARTUP_NOTICE_SECONDS = 15.0 +TERMINAL_OUTPUT_QUEUE_LIMIT = 256 EXTRACT_PROVIDERS = {"claude", "codex", "cursor", "hermes", "pi"} UPLOAD_IGNORE_PATTERNS = ["partials/**", "failures/**"] UPLOAD_METADATA_PATTERNS = ["README.md", "tools.json"] +UPLOAD_DATA_PATTERNS = ["*.jsonl", "**/*.jsonl", "*.metadata.json", "**/*.metadata.json"] def _normalize_extract_provider(provider: str) -> ExtractProvider: @@ -241,7 +243,7 @@ def stream(): if events: for event in events: yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n" - index += len(events) + index = int(events[-1]["seq"]) + 1 idle_cycles = 0 else: if log.closed: @@ -331,6 +333,7 @@ def _upload_dataset_folder( folder_path=str(folder_path), repo_type="dataset", private=private, + allow_patterns=UPLOAD_DATA_PATTERNS, ignore_patterns=list(dict.fromkeys([*ignore_patterns, *UPLOAD_METADATA_PATTERNS])), ) else: @@ -339,6 +342,7 @@ def _upload_dataset_folder( repo_id=repo_id, repo_type="dataset", commit_message="Upload teich dataset output", + allow_patterns=[*UPLOAD_DATA_PATTERNS, *UPLOAD_METADATA_PATTERNS], ignore_patterns=ignore_patterns, ) return str(repo_url) @@ -519,10 +523,10 @@ def start_extraction(payload: ExtractRequest) -> dict[str, Any]: ] model_filter = (payload.model or "").strip() or None config_output = str(output_dir) if Path(output_value).expanduser().is_absolute() else output_value - try: + + def persist_output_config() -> None: state.write_config_data({"output": {"traces_dir": config_output}}) - except Exception: - pass + try: job = extraction.start( provider, @@ -530,9 +534,12 @@ def start_extraction(payload: ExtractRequest) -> dict[str, Any]: source_paths=source_paths, model_filter=model_filter, skip_anonymize=payload.skip_anonymize, + before_start=persist_output_config, ) except RuntimeError as exc: raise HTTPException(status_code=409, detail=str(exc)) + except Exception as exc: + raise HTTPException(status_code=400, detail=str(exc)) return job.to_dict() @app.get("/api/extract") @@ -635,10 +642,18 @@ async def session_terminal(websocket: WebSocket, session_id: str, cols: int = 12 return await websocket.accept() loop = asyncio.get_running_loop() - queue: asyncio.Queue[str] = asyncio.Queue() + queue: asyncio.Queue[str] = asyncio.Queue(maxsize=TERMINAL_OUTPUT_QUEUE_LIMIT) + + def enqueue_output(text: str) -> None: + if queue.full(): + try: + queue.get_nowait() + except asyncio.QueueEmpty: + pass + queue.put_nowait(text) def on_output(text: str) -> None: - loop.call_soon_threadsafe(queue.put_nowait, text) + loop.call_soon_threadsafe(enqueue_output, text) await _wait_for_terminal_session_ready(websocket, session) if session.status == "error": diff --git a/tests/test_audit.py b/tests/test_audit.py index cee7126..40661ab 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -97,6 +97,17 @@ def test_audit_sft_dataset_rejects_gemma_context_markers(): assert "<|tool_response>" in report.errors[1] +def test_audit_sft_dataset_checks_all_rows_by_default(): + safe = {"input_ids": [1, 2, 4], "attention_mask": [1, 1, 1], "labels": [-100, 2, 4]} + leaked = {"input_ids": [6, 2], "attention_mask": [1, 1], "labels": [6, 2]} + dataset = Dataset.from_list([safe] * 8 + [leaked]) + + report = audit_sft_dataset(dataset, TinyTokenizer()) + + assert not report.ok + assert any("row 8" in error and "<|turn>user" in error for error in report.errors) + + def test_teich_example_has_single_safe_training_flow(): source = Path("teich_example.py").read_text(encoding="utf-8") tree = ast.parse(source) diff --git a/tests/test_cli.py b/tests/test_cli.py index b669bcb..b1e53e7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -513,9 +513,10 @@ def test_generate_command_publishes_dataset_when_publish_repo_is_configured(tmp_ mock_api.upload_large_folder.assert_called_once_with( repo_id="armand0e/test-dataset", folder_path=str(output_dir), - repo_type="dataset", - private=True, - ignore_patterns=["partials/**", "failures/**", "README.md", "tools.json"], + repo_type="dataset", + private=True, + allow_patterns=["*.jsonl", "**/*.jsonl", "*.metadata.json", "**/*.metadata.json"], + ignore_patterns=["partials/**", "failures/**", "README.md", "tools.json"], ) mock_api.upload_folder.assert_called_once_with( folder_path=str(output_dir), @@ -669,9 +670,10 @@ def test_generate_command_can_publish_completed_outputs_after_failure(tmp_path: mock_api.upload_large_folder.assert_called_once_with( repo_id="armand0e/test-dataset", folder_path=str(output_dir), - repo_type="dataset", - private=False, - ignore_patterns=["partials/**", "failures/**", "README.md", "tools.json"], + repo_type="dataset", + private=False, + allow_patterns=["*.jsonl", "**/*.jsonl", "*.metadata.json", "**/*.metadata.json"], + ignore_patterns=["partials/**", "failures/**", "README.md", "tools.json"], ) mock_api.upload_folder.assert_called_once_with( folder_path=str(output_dir), diff --git a/tests/test_config.py b/tests/test_config.py index d0be29d..4910858 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -34,6 +34,13 @@ def test_default_config(): ] +def test_config_rejects_unknown_top_level_and_nested_fields(): + with pytest.raises(ValueError, match="max_concurency"): + Config.model_validate({"max_concurency": 99}) + with pytest.raises(ValueError, match="reasoning_efort"): + Config.model_validate({"model": {"reasoning_efort": "high"}}) + + @pytest.mark.parametrize("field", ["startup_timeout_sec", "tool_timeout_sec"]) def test_mcp_timeouts_must_be_positive(field: str): with pytest.raises(ValueError): diff --git a/tests/test_extract_anonymize_cli.py b/tests/test_extract_anonymize_cli.py index 7f88cb5..a88ca62 100644 --- a/tests/test_extract_anonymize_cli.py +++ b/tests/test_extract_anonymize_cli.py @@ -1,3 +1,4 @@ +import base64 import json import os import sqlite3 @@ -1694,6 +1695,7 @@ def test_extract_can_upload_staged_anonymized_output_to_huggingface(tmp_path: Pa folder_path=str(output_dir), repo_type="dataset", private=False, + allow_patterns=["*.jsonl", "**/*.jsonl", "*.metadata.json", "**/*.metadata.json"], ignore_patterns=["partials/**", "failures/**", "README.md", "tools.json"], ) mock_api.upload_folder.assert_called_once_with( @@ -2015,8 +2017,7 @@ def test_anonymize_generalizes_to_synthetic_secret_and_reference_matrix(tmp_path assert "PROJECT_SECRET=redacted_" in text assert "process.env.PROJECT_SECRET" in text assert "sk-Mmu5OJR5NQoTJOGj6z6cHPraZ35yuZfK" not in text - assert image_data in text - assert "redacted_base64" not in text + assert image_data not in text assert "@src/main.ts" in text assert "@docs/*.md" in text assert "@scope/package" in text @@ -2210,7 +2211,7 @@ def test_anonymize_scrubs_env_style_keys_without_scrubbing_tokenizer_terms(tmp_p assert "credential_occurrences=3" in result.output -def test_anonymize_preserves_base64_media_payloads_without_touching_metadata(tmp_path: Path): +def test_anonymize_redacts_base64_media_payloads_without_touching_metadata(tmp_path: Path): input_dir = tmp_path / "output" input_dir.mkdir() image_data = "aGVsbG8=" * 40 @@ -2241,9 +2242,9 @@ def test_anonymize_preserves_base64_media_payloads_without_touching_metadata(tmp source = row["content"][0]["source"] assert source["type"] == "base64" assert source["media_type"] == "image/png" - assert source["data"] == image_data - assert image_data in json.dumps(row) - assert "image_data" not in result.output + assert source["data"] != image_data + assert image_data not in json.dumps(row) + assert base64.b64decode(source["data"]).startswith(b"[redacted media ") def test_anonymize_does_not_treat_systemd_units_as_emails(tmp_path: Path): @@ -2349,7 +2350,7 @@ def test_anonymize_keeps_high_confidence_credentials_and_signed_query_values(): assert anonymizer.counts["api_key"] == 3 -def test_anonymize_preserves_assistant_authored_examples_but_scans_tool_inputs(): +def test_anonymize_scans_assistant_authored_examples_and_tool_inputs(): sample_key = "sk-proj-abcdefghijklmnopqrstuvwxyzABCDEFGHIJK" sample_jwt = ( "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9." @@ -2385,14 +2386,55 @@ def test_anonymize_preserves_assistant_authored_examples_but_scans_tool_inputs() redacted = anonymizer.anonymize_value(value) assistant = redacted["messages"][0] - assert sample_key in assistant["content"] - assert sample_jwt in assistant["content"] - assert "correct-horse-battery-staple" in assistant["content"] - assert "developer@company.com" in assistant["content"] - assert "qa@acme.com" in assistant["reasoning_content"] + assert sample_key not in assistant["content"] + assert sample_jwt not in assistant["content"] + assert "correct-horse-battery-staple" not in assistant["content"] + assert "developer@company.com" not in assistant["content"] + assert "qa@acme.com" not in assistant["reasoning_content"] assert tool_secret not in json.dumps(assistant["tool_calls"]) assert "redacted_secret_" in json.dumps(assistant["tool_calls"]) - assert anonymizer.counts == {"email": 0, "username": 0, "api_key": 1} + assert anonymizer.counts == {"email": 2, "username": 0, "api_key": 4, "pii": 0, "media": 0} + + +def test_anonymize_scrubs_high_confidence_personal_identifiers(): + value = { + "messages": [ + { + "role": "assistant", + "content": ( + "My name is Jane Example. Phone: (212) 555-0198. " + "Address: 123 Example Street, New York, NY 10001. " + "SSN 123-45-6789. IP address: 192.0.2.45." + ), + } + ] + } + + anonymizer = anonymize_module.TraceAnonymizer() + redacted = anonymizer.anonymize_value(value) + text = redacted["messages"][0]["content"] + + for private_value in ( + "Jane Example", + "(212) 555-0198", + "123 Example Street", + "123-45-6789", + "192.0.2.45", + ): + assert private_value not in text + assert text.count("redacted_pii_") == 5 + assert anonymizer.counts["pii"] == 5 + + structured = anonymizer.anonymize_value( + { + "full_name": "Sam Private", + "phone_number": "212-555-0114", + "street_address": "99 Private Road", + "ip_address": "198.51.100.7", + } + ) + assert all(str(value).startswith("redacted_pii_") for value in structured.values()) + assert anonymizer.counts["pii"] == 9 def test_anonymize_generic_secret_scan_is_bounded_on_long_hyphenated_ids(): diff --git a/tests/test_formatter.py b/tests/test_formatter.py index eee821d..e09fa4b 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -449,6 +449,126 @@ def apply_chat_template( return rendered +def test_gemma4_rejects_reasoning_when_thinking_is_disabled(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-26B-A4B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer", "reasoning_content": "reason"}, + ], + "tools": [], + } + ] + ) + + with pytest.raises(ValueError, match="Gemma 4 thinking is disabled"): + prepare_data( + dataset, + tokenizer, + chat_template_kwargs={"enable_thinking": False}, + strict=True, + verbose=False, + ) + + +@pytest.mark.parametrize("enable_thinking", [False, True]) +def test_gemma4_rejects_manual_think_trigger_in_source_data(enable_thinking: bool): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-E4B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "system", "content": "<|think|>"}, + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ], + "tools": [], + } + ] + ) + + with pytest.raises(ValueError, match="Do not embed the trigger manually"): + prepare_data( + dataset, + tokenizer, + chat_template_kwargs={ + "enable_thinking": enable_thinking, + "preserve_thinking": enable_thinking, + }, + strict=True, + verbose=False, + ) + + +def test_gemma4_rejects_dropped_historical_reasoning(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-31B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "one", "reasoning_content": "reason one"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "two", "reasoning_content": "reason two"}, + ], + "tools": [], + } + ] + ) + + with pytest.raises(ValueError, match="preserve_thinking=False"): + prepare_data( + dataset, + tokenizer, + chat_template_kwargs={"enable_thinking": True, "preserve_thinking": False}, + strict=True, + verbose=False, + ) + + +def test_gemma4_nonthinking_training_matches_empty_thought_generation_prefix(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-26B-A4B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ], + "tools": [], + } + ] + ) + + prepared = prepare_data( + dataset, + tokenizer, + chat_template_kwargs={"enable_thinking": False}, + strict=True, + verbose=False, + ) + + assert "<|turn>model\n<|channel>thought\nanswer" in prepared[0]["text"] + training_data = prepare_and_mask_for_test( + dataset, + tokenizer, + chat_template_kwargs={"enable_thinking": False}, + strict=True, + ) + supervised_text = tokenizer.decode( + [token for token in training_data[0]["labels"] if token != -100] + ) + assert "<|channel>thought" not in supervised_text + assert "answer" in supervised_text + assert supervised_text.endswith("") + + def test_reordered_mapping_markers_preserve_typed_tool_spans(): class SortedToolArgumentTokenizer(OffsetCountingTokenizer): def apply_chat_template(self, messages, *, tokenize=False, add_generation_prompt=False, tools=None, **kwargs): diff --git a/tests/test_runner.py b/tests/test_runner.py index 959babb..53e0a38 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -5603,7 +5603,10 @@ def test_resume_treats_prompt_level_system_as_part_of_chat_completion_key(tmp_pa pending = pending_prompt_inputs_for_resume(prompt_inputs, output_dir) - assert pending == [] + assert [(item.system, item.prompt) for item in pending] == [ + (None, "Hello"), + ("Be thorough.", "Hello"), + ] def test_resume_matches_completed_agent_trace_when_configured_prompt_has_system(tmp_path: Path): @@ -5624,7 +5627,7 @@ def test_resume_matches_completed_agent_trace_when_configured_prompt_has_system( pending = pending_prompt_inputs_for_resume(prompt_inputs, output_dir) - assert pending == [] + assert pending == prompt_inputs def test_resume_detects_completed_chat_follow_up_prompt_sets(tmp_path: Path): @@ -6655,7 +6658,7 @@ def test_pi_runner_appends_prompt_level_system_metadata_at_end(tmp_path: Path): ] trace_file.write_text("\n".join(native_lines) + "\n", encoding="utf-8") - runner._append_pi_system_prompt_metadata( + runner._append_prompt_system_metadata( trace_file, PromptInput(prompt="Hello", system="Use the prompt-level system."), ) @@ -6754,7 +6757,7 @@ def test_pi_runner_does_not_duplicate_prompt_level_system_metadata(tmp_path: Pat ) before = trace_file.read_text(encoding="utf-8") - runner._append_pi_system_prompt_metadata( + runner._append_prompt_system_metadata( trace_file, PromptInput(prompt="Hello", system="Use the prompt-level system."), ) diff --git a/tests/test_studio.py b/tests/test_studio.py index abeed16..6df7a97 100644 --- a/tests/test_studio.py +++ b/tests/test_studio.py @@ -18,10 +18,13 @@ from teich.extract import CURSOR_EXTRACTION_NOTICE from teich.runner import ChatRunner, ClaudeCodeRunner, SessionProgressUpdate from teich.studio.events import summarize_chat_row, summarize_event, summarize_trace_events +from teich.studio.extraction import ExtractionJob, ExtractionManager from teich.studio.generation import RUNNER_CLASSES, GenerationJob from teich.studio.interactive import ( SUBPROCESS_CREATE_NO_WINDOW, + EventLog, InteractiveSession, + SessionManager, TerminalBridge, ) from teich.studio.project import ProjectState @@ -72,6 +75,31 @@ def test_write_config_rejects_invalid(tmp_path): state.write_config_data({"max_concurrency": 0}) +def test_event_log_retains_a_bounded_absolute_sequence(): + log = EventLog(max_events=3) + for index in range(5): + log.append({"kind": "test", "value": index}) + + assert [event["seq"] for event in log.snapshot()] == [2, 3, 4] + assert [event["value"] for event in log.wait_for(0, timeout=0)] == [2, 3, 4] + assert [event["value"] for event in log.wait_for(4, timeout=0)] == [4] + + +def test_session_manager_prunes_old_completed_history(): + manager = SessionManager() + manager._sessions = { + str(index): SimpleNamespace(status="saved") + for index in range(55) + } + + with manager._lock: + manager._prune_locked() + + assert len(manager._sessions) == 49 + assert "0" not in manager._sessions + assert "54" in manager._sessions + + def test_prompts_round_trip(tmp_path): state = ProjectState(tmp_path) state.ensure_initialized() @@ -542,6 +570,7 @@ def test_dataset_upload_endpoint_generates_readme_and_uploads_with_env_token(cli folder_path=str(output_dir), repo_type="dataset", private=False, + allow_patterns=["*.jsonl", "**/*.jsonl", "*.metadata.json", "**/*.metadata.json"], ignore_patterns=["partials/**", "failures/**", "README.md", "tools.json"], ) mock_api.upload_folder.assert_called_once_with( @@ -861,6 +890,46 @@ def test_extract_endpoint_can_skip_anonymization(client): assert "raw@company.ai" in (client.project_dir / "raw-staged" / "raw.jsonl").read_text(encoding="utf-8") +def test_rejected_extract_does_not_mutate_output_config(client): + before = client.get("/api/config").json()["config"]["output"]["traces_dir"] + client.app.state.extraction._current = SimpleNamespace(status="running", join=lambda: None) + + response = client.post( + "/api/extract", + json={"provider": "codex", "output": "rejected-output", "sessions_dirs": []}, + ) + + assert response.status_code == 409 + after = client.get("/api/config").json()["config"]["output"]["traces_dir"] + assert after == before + + +def test_extraction_shutdown_waits_for_active_worker(tmp_path): + started = threading.Event() + release = threading.Event() + shutdown_finished = threading.Event() + + def blocked_run(job): + job.status = "running" + started.set() + release.wait(timeout=5) + job.status = "completed" + + manager = ExtractionManager() + with patch.object(ExtractionJob, "_run", blocked_run): + manager.start("codex", output_dir=tmp_path) + assert started.wait(timeout=1) + thread = threading.Thread( + target=lambda: (manager.shutdown(), shutdown_finished.set()), + daemon=True, + ) + thread.start() + assert not shutdown_finished.wait(timeout=0.05) + release.set() + assert shutdown_finished.wait(timeout=1) + thread.join(timeout=1) + + def test_extract_endpoint_warns_cursor_may_take_a_while(client): source = client.project_dir / "state.vscdb" source.write_text("", encoding="utf-8") diff --git a/tests/test_tokenizer_smoke.py b/tests/test_tokenizer_smoke.py index 2d30f0b..ccb7e9a 100644 --- a/tests/test_tokenizer_smoke.py +++ b/tests/test_tokenizer_smoke.py @@ -11,9 +11,21 @@ TOKENIZER_SMOKE_MODELS = [ pytest.param("unsloth/Qwen3.5-0.8B", {"enable_thinking": True}, id="unsloth-qwen3.5"), - pytest.param("google/gemma-4-E4B-it", {"enable_thinking": True}, id="gemma-4-e4b-it"), - pytest.param("google/gemma-4-26B-A4B-it", {"enable_thinking": True}, id="gemma-4-26b-a4b-it"), - pytest.param("google/gemma-4-31B-it", {"enable_thinking": True}, id="gemma-4-31b-it"), + pytest.param( + "google/gemma-4-E4B-it", + {"enable_thinking": True, "preserve_thinking": True}, + id="gemma-4-e4b-it", + ), + pytest.param( + "google/gemma-4-26B-A4B-it", + {"enable_thinking": True, "preserve_thinking": True}, + id="gemma-4-26b-a4b-it", + ), + pytest.param( + "google/gemma-4-31B-it", + {"enable_thinking": True, "preserve_thinking": True}, + id="gemma-4-31b-it", + ), ] @@ -111,6 +123,7 @@ def test_real_tokenizer_prepare_and_mask_tool_dataset(model_id: str, chat_templa assert "SECRET_TOOL_OUTPUT" in masked_text if model_id.startswith("google/gemma-4-"): + assert "<|think|>" in prepared[0]["text"] eot_token_id = tokenizer.convert_tokens_to_ids("") assert eot_token_id in supervised_ids assert supervised_text.endswith("") @@ -125,3 +138,47 @@ def test_real_tokenizer_prepare_and_mask_tool_dataset(model_id: str, chat_templa **chat_template_kwargs, ) assert generation_prompt.endswith("<|channel>thought\n") + + +@pytest.mark.integration +@pytest.mark.tokenizer_smoke +@pytest.mark.parametrize( + "model_id,expects_empty_thought", + [ + pytest.param("google/gemma-4-E4B-it", False, id="gemma-4-e4b-direct"), + pytest.param("google/gemma-4-26B-A4B-it", True, id="gemma-4-26b-a4b-direct"), + pytest.param("google/gemma-4-31B-it", True, id="gemma-4-31b-direct"), + ], +) +def test_real_gemma4_nonthinking_training_matches_generation_prompt( + model_id: str, + expects_empty_thought: bool, +): + if not _tokenizer_smokes_enabled(): + pytest.skip("Set TEICH_RUN_TOKENIZER_SMOKES=1 to run real Hugging Face tokenizer smokes.") + transformers = pytest.importorskip("transformers") + tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + messages = [ + {"role": "user", "content": "Give a direct answer."}, + {"role": "assistant", "content": "Direct answer."}, + ] + + prepared = prepare_data( + Dataset.from_list([{"messages": messages, "tools": []}]), + tokenizer, + tokenize=True, + strict=True, + chat_template_kwargs={"enable_thinking": False}, + verbose=False, + ) + empty_thought = "<|channel>thought\n" + generation_prompt = tokenizer.apply_chat_template( + messages[:-1], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + + assert (empty_thought in generation_prompt) is expects_empty_thought + assert (empty_thought in prepared[0]["text"]) is expects_empty_thought + assert "<|think|>" not in prepared[0]["text"] From 57ccdf00f35e9549876bf1b5e08d57508c92c203 Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Mon, 24 Aug 2026 23:32:09 -0400 Subject: [PATCH 2/6] Automate Gemma 4 thinking mode selection --- README.md | 4 + docs/prepare-data.md | 10 +- docs/training.md | 78 +++++++-------- gemma4_example.py | 26 +++-- src/teich/anonymize.py | 18 +++- src/teich/formatter.py | 141 ++++++++++++++++++++++++---- tests/test_audit.py | 4 + tests/test_extract_anonymize_cli.py | 27 ++++-- tests/test_formatter.py | 122 ++++++++++++++++++++++-- tests/test_tokenizer_smoke.py | 27 +++++- 10 files changed, 375 insertions(+), 82 deletions(-) diff --git a/README.md b/README.md index a26d32b..6b7f5c7 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,10 @@ train_dataset = prepare_data( ) ``` +With a live Gemma 4 tokenizer, omit `chat_template_kwargs` to let Teich choose +thinking or non-thinking independently for each row. See [Training](docs/training.md#live-gemma-4-models) +for auto-mode rules and the E4B, 26B-A4B, and 31B template contract. + Then create your trainer and call `mask_data()`: ```python diff --git a/docs/prepare-data.md b/docs/prepare-data.md index 3a8ec4f..0053f37 100644 --- a/docs/prepare-data.md +++ b/docs/prepare-data.md @@ -48,7 +48,15 @@ print(prep_report.max_token_length) print(prep_report.oversized_rows[:3]) ``` -`PrepareReport` includes dropped rows, oversized rows, trimmed rows, token lengths, max token length, kept-row ids, and returned row count. +`PrepareReport` includes dropped rows, oversized rows, trimmed rows, token lengths, +max token length, kept-row ids, and returned row count. With a live Gemma 4 +template it also reports per-row mode counts in `gemma4_modes` and migrated +leading `<|think|>` markers in `gemma4_legacy_triggers_normalized`. + +Gemma 4 defaults to per-row auto mode when `enable_thinking` is omitted: +reasoning-bearing rows enable thinking and history preservation, while direct +rows use the loaded model's non-thinking protocol. See [Live Gemma 4 Models](training.md#live-gemma-4-models) +for explicit overrides and validation rules. Original columns are removed after formatting unless `preserve_columns=True` or an explicit list is passed. The default provenance set is `source`, `metadata`, `raw_index`, and `source_key`. diff --git a/docs/training.md b/docs/training.md index bf2ea9c..c63b0e4 100644 --- a/docs/training.md +++ b/docs/training.md @@ -57,58 +57,62 @@ are: - `google/gemma-4-26B-A4B-it` - `google/gemma-4-31B-it` -Gemma 4 does not infer this mode from the presence of reasoning labels alone. -The live template enables thinking by placing `<|think|>` at the beginning of -the system turn when `enable_thinking=True`; do not embed that token manually, -because enabling the template would then duplicate it. For thinking SFT, pass -`chat_template_kwargs={"enable_thinking": True, "preserve_thinking": True}`. -`preserve_thinking=True` is required for multi-turn rows with historical -reasoning; otherwise the upstream template silently omits reasoning from all -but the last assistant turn. - -For non-thinking SFT, pass `{"enable_thinking": False}` and use rows that have -neither reasoning fields nor a manually embedded `<|think|>` system trigger. -The live template will render `reasoning_content` even when thinking is off, -which creates an internally inconsistent example; Teich rejects that state. -The 26B-A4B and 31B generation prompts also insert an empty thought channel in -non-thinking mode while E4B does not. Teich detects this from the loaded live -template and mirrors the prefix in completed training turns so SFT and -inference contexts agree. +Gemma 4's thinking mode is a prompt protocol, so Teich resolves it for every +row before rendering. Leave `enable_thinking` unset for the recommended auto +mode: + +- a row containing assistant `reasoning_content`, `thinking`, or `reasoning` + becomes a thinking row; +- a row without assistant reasoning becomes a non-thinking row; +- thinking rows automatically receive `enable_thinking=True` and + `preserve_thinking=True`, preserving historical reasoning in multi-turn data; +- a leading legacy `<|think|>` in system/developer content is removed and used + as a thinking-mode hint. `PrepareReport` records the migration, while triggers + in any other position are rejected. + +The live template, not the source data, inserts `<|think|>` into the system +turn. The 26B-A4B and 31B generation prompts also insert an empty thought +channel in non-thinking mode while E4B does not. Teich discovers that behavior +from the loaded template and mirrors the prefix in completed SFT turns while +keeping the synthetic prefix masked. + +Explicit `{"enable_thinking": True}` and `{"enable_thinking": False}` remain +available when an entire source must be forced to one mode. Teich automatically +enables history preservation for the explicit thinking case. Contradictions +fail closed: forcing non-thinking on a reasoning-bearing row, or explicitly +disabling `preserve_thinking` when historical reasoning would disappear, +raises an error instead of producing inconsistent training text. Do not replace `tokenizer.chat_template` unless you are intentionally testing a maintained fork. Teich supervises the closing `` token for completed Gemma responses while keeping system, user, tool-response, and generated prompt prefix context masked. -Thinking and non-thinking examples can be mixed safely by separating them into -sources with source-level template kwargs: +Thinking and non-thinking examples can therefore be mixed in the same source +without any template configuration: ```python -train_dataset = prepare_data( - { - "thinking": { - "source": "username/gemma-thinking-traces", - "percentage": 60, - "chat_template_kwargs": { - "enable_thinking": True, - "preserve_thinking": True, - }, - }, - "direct": { - "source": "username/gemma-direct-traces", - "percentage": 40, - "chat_template_kwargs": {"enable_thinking": False}, - }, - }, +train_dataset, prep_report = prepare_data( + "username/mixed-gemma-traces", tokenizer, tokenize=True, strict=True, + return_report=True, ) ``` +With `return_report=True`, inspect `prep_report.gemma4_modes` for the resolved +thinking/non-thinking counts and +`prep_report.gemma4_legacy_triggers_normalized` for migrated legacy rows. +Source-level `chat_template_kwargs` are still useful as explicit policy +overrides, but are no longer required merely to mix the two modes. + `gemma4_example.py` uses the live remote template by default. Set -`CHAT_TEMPLATE_PATH` only to opt into a local custom template, and set -`MODEL_REVISION` when a reproducible non-`main` revision is required. Set +`CHAT_TEMPLATE_PATH` only to opt into a local custom template. Its +`GEMMA4_THINKING_MODE` defaults to `auto`; use `thinking` or `nonthinking` only +to force a homogeneous source. The older `GEMMA4_ENABLE_THINKING` variable is +still accepted for compatibility. Set `MODEL_REVISION` when a reproducible +non-`main` revision is required. Set `HF_TOKEN` to an account that has accepted the Gemma repository terms when the checkpoint is not already available through the local Hugging Face login. diff --git a/gemma4_example.py b/gemma4_example.py index 427663d..3061da5 100644 --- a/gemma4_example.py +++ b/gemma4_example.py @@ -17,7 +17,21 @@ HUB_REPO_ID = os.environ.get("HUB_REPO_ID") or "" HF_TOKEN = os.environ.get("HF_TOKEN", "") CHAT_TEMPLATE_PATH = os.environ.get("CHAT_TEMPLATE_PATH") -ENABLE_THINKING = os.environ.get("GEMMA4_ENABLE_THINKING", "1").strip().lower() not in {"0", "false", "no"} +_thinking_mode = os.environ.get("GEMMA4_THINKING_MODE") +_legacy_thinking = os.environ.get("GEMMA4_ENABLE_THINKING") +if _thinking_mode is None and _legacy_thinking is not None: + _thinking_mode = ( + "nonthinking" + if _legacy_thinking.strip().lower() in {"0", "false", "no"} + else "thinking" + ) +GEMMA4_THINKING_MODE = (_thinking_mode or "auto").strip().lower().replace("-", "") +if GEMMA4_THINKING_MODE not in {"auto", "thinking", "nonthinking"}: + raise ValueError("GEMMA4_THINKING_MODE must be auto, thinking, or nonthinking") +CHAT_TEMPLATE_KWARGS = { + "thinking": {"enable_thinking": True}, + "nonthinking": {"enable_thinking": False}, +}.get(GEMMA4_THINKING_MODE) model, tokenizer = FastModel.from_pretrained( model_name=MODEL_NAME, @@ -67,12 +81,10 @@ tokenizer, split="train", hf_token=HF_TOKEN, - # Non-thinking datasets must not contain reasoning fields or a manual - # <|think|> system trigger; Teich rejects those inconsistent rows. - chat_template_kwargs={ - "enable_thinking": ENABLE_THINKING, - "preserve_thinking": ENABLE_THINKING, - }, + # Auto mode classifies every Gemma 4 row independently. Reasoning-bearing + # rows enable thinking and preserve history; direct rows use the exact + # non-thinking inference prefix of the loaded live template. + chat_template_kwargs=CHAT_TEMPLATE_KWARGS, max_length=MAX_SEQ_LEN, oversized_policy="trim_followups", tokenize=True, diff --git a/src/teich/anonymize.py b/src/teich/anonymize.py index 217fe3b..3359927 100644 --- a/src/teich/anonymize.py +++ b/src/teich/anonymize.py @@ -36,6 +36,7 @@ # ponytail: process startup isn't free — only fan out when there are enough # files for the parallelism to pay for itself. _MIN_FILES_FOR_PARALLEL = 8 +_MIN_TOTAL_BYTES_FOR_PARALLEL = 1024 * 1024 _WINDOWS_MAX_PROCESS_WORKERS = 61 @@ -46,6 +47,21 @@ def _process_worker_count(file_count: int) -> int: return workers +def _use_process_workers(source_files: list[Path], workers: int) -> bool: + if workers <= 1 or len(source_files) < _MIN_FILES_FOR_PARALLEL: + return False + total_bytes = 0 + for source_file in source_files: + try: + total_bytes += source_file.stat().st_size + except OSError: + # Let anonymize_file surface the underlying filesystem error. + return True + if total_bytes >= _MIN_TOTAL_BYTES_FOR_PARALLEL: + return True + return False + + @dataclass class AnonymizeFileReport: path: Path @@ -119,7 +135,7 @@ def anonymize_files( raise ValueError("source_files and destinations must have the same length") reports: list[AnonymizeFileReport] = [] workers = _process_worker_count(len(source_files)) - if workers > 1 and len(source_files) >= _MIN_FILES_FOR_PARALLEL: + if _use_process_workers(source_files, workers): # Each file is anonymized independently (fresh TraceAnonymizer per # file), so files can be processed in parallel safely. # Extraction can run from a Studio background thread. Explicit spawn diff --git a/src/teich/formatter.py b/src/teich/formatter.py index f8c70df..4f6af58 100644 --- a/src/teich/formatter.py +++ b/src/teich/formatter.py @@ -98,6 +98,8 @@ class PrepareReport: dropped_rows: list[dict[str, Any]] = field(default_factory=list) oversized_rows: list[dict[str, Any]] = field(default_factory=list) trimmed_rows: list[dict[str, Any]] = field(default_factory=list) + gemma4_modes: dict[str, int] = field(default_factory=dict) + gemma4_legacy_triggers_normalized: int = 0 def record_token_length(self, row_info: dict[str, Any], token_length: int) -> None: entry = {**row_info, "token_length": token_length} @@ -153,6 +155,13 @@ def record_trimmed_row( } ) + def record_gemma4_mode(self, mode: str | None, *, normalized_legacy_trigger: bool) -> None: + if mode is None: + return + self.gemma4_modes[mode] = self.gemma4_modes.get(mode, 0) + 1 + if normalized_legacy_trigger: + self.gemma4_legacy_triggers_normalized += 1 + def to_dict(self) -> dict[str, Any]: return { "total_rows": self.total_rows, @@ -165,6 +174,8 @@ def to_dict(self) -> dict[str, Any]: "dropped_rows": self.dropped_rows, "oversized_rows": self.oversized_rows, "trimmed_rows": self.trimmed_rows, + "gemma4_modes": self.gemma4_modes, + "gemma4_legacy_triggers_normalized": self.gemma4_legacy_triggers_normalized, } @@ -174,6 +185,8 @@ class _RenderedRow: supervised_spans: list[dict[str, Any]] tokenized: tuple[list[int], list[int]] | None token_length: int | None + gemma4_mode: str | None = None + normalized_legacy_trigger: bool = False @dataclass(slots=True) @@ -253,41 +266,106 @@ def _message_text_contains(message: dict[str, Any], needle: str) -> bool: return False -def _validate_gemma4_thinking_contract( +def _strip_leading_gemma4_trigger(message: dict[str, Any]) -> tuple[dict[str, Any], bool]: + content = message.get("content") + if isinstance(content, str): + stripped = content.lstrip() + if not stripped.startswith(_GEMMA_THINK_TRIGGER): + return message, False + updated = dict(message) + updated["content"] = stripped[len(_GEMMA_THINK_TRIGGER):].lstrip() + return updated, True + if not isinstance(content, list): + return message, False + updated_content = deepcopy(content) + for index, part in enumerate(updated_content): + if isinstance(part, str): + if not part.strip(): + continue + stripped = part.lstrip() + if not stripped.startswith(_GEMMA_THINK_TRIGGER): + return message, False + updated_content[index] = stripped[len(_GEMMA_THINK_TRIGGER):].lstrip() + updated = dict(message) + updated["content"] = updated_content + return updated, True + if not isinstance(part, dict): + return message, False + text = part.get("text") + if not isinstance(text, str) or not text.strip(): + continue + stripped = text.lstrip() + if not stripped.startswith(_GEMMA_THINK_TRIGGER): + return message, False + updated_part = dict(part) + updated_part["text"] = stripped[len(_GEMMA_THINK_TRIGGER):].lstrip() + updated_content[index] = updated_part + updated = dict(message) + updated["content"] = updated_content + return updated, True + return message, False + + +def _resolve_gemma4_thinking_contract( renderer: Any, messages: list[dict[str, Any]], chat_template_kwargs: dict[str, Any], -) -> None: +) -> tuple[list[dict[str, Any]], dict[str, Any], str | None, bool]: if not _is_gemma4_renderer(renderer): - return - thinking_enabled = chat_template_kwargs.get("enable_thinking", False) is True + return messages, chat_template_kwargs, None, False + resolved_messages = messages + normalized_legacy_trigger = False + for index, message in enumerate(messages): + if message.get("role") not in {"system", "developer"}: + continue + if not _message_text_contains(message, _GEMMA_THINK_TRIGGER): + continue + updated_message, stripped = _strip_leading_gemma4_trigger(message) + if not stripped or _message_text_contains(updated_message, _GEMMA_THINK_TRIGGER): + raise ValueError( + "Gemma 4 system/developer content contains <|think|> outside the supported leading " + "legacy position. Remove it and let Teich manage the thinking trigger." + ) + if resolved_messages is messages: + resolved_messages = list(messages) + resolved_messages[index] = updated_message + normalized_legacy_trigger = True + reasoning_indexes = [ index - for index, message in enumerate(messages) + for index, message in enumerate(resolved_messages) if message.get("role") in {"assistant", "model"} and _message_reasoning(message) ] - manual_trigger = any( - message.get("role") in {"system", "developer"} - and _message_text_contains(message, _GEMMA_THINK_TRIGGER) - for message in messages + resolved_kwargs = dict(chat_template_kwargs) + explicit_thinking = "enable_thinking" in resolved_kwargs + if explicit_thinking and not isinstance(resolved_kwargs["enable_thinking"], bool): + raise ValueError("Gemma 4 enable_thinking must be a boolean when provided.") + if "preserve_thinking" in resolved_kwargs and not isinstance(resolved_kwargs["preserve_thinking"], bool): + raise ValueError("Gemma 4 preserve_thinking must be a boolean when provided.") + thinking_enabled = ( + resolved_kwargs["enable_thinking"] + if explicit_thinking + else bool(reasoning_indexes or normalized_legacy_trigger) ) - if manual_trigger: + if explicit_thinking and not thinking_enabled and normalized_legacy_trigger: raise ValueError( - "Gemma 4 system/developer content already contains <|think|>. Do not embed the trigger " - "manually: use chat_template_kwargs={'enable_thinking': True, " - "'preserve_thinking': True} so the live template inserts it exactly once." + "Gemma 4 enable_thinking=False conflicts with a leading <|think|> system trigger. " + "Remove the legacy trigger for non-thinking data or omit enable_thinking for auto mode." ) if not thinking_enabled and reasoning_indexes: raise ValueError( "Gemma 4 thinking is disabled, but this row contains reasoning. Use " "chat_template_kwargs={'enable_thinking': True, " - "'preserve_thinking': True} for thinking data, or remove reasoning for non-thinking data." + "'preserve_thinking': True}, omit enable_thinking for auto mode, " + "or remove reasoning for non-thinking data." ) - if thinking_enabled and chat_template_kwargs.get("preserve_thinking", False) is not True: + resolved_kwargs["enable_thinking"] = thinking_enabled + resolved_kwargs.setdefault("preserve_thinking", thinking_enabled) + if thinking_enabled and resolved_kwargs["preserve_thinking"] is not True: last_assistant = max( ( index - for index, message in enumerate(messages) + for index, message in enumerate(resolved_messages) if message.get("role") in {"assistant", "model"} ), default=-1, @@ -297,6 +375,8 @@ def _validate_gemma4_thinking_contract( "Gemma 4 preserve_thinking=False drops reasoning from earlier assistant turns. " "Set preserve_thinking=True for multi-turn thinking SFT." ) + mode = "thinking" if thinking_enabled else "nonthinking" + return resolved_messages, resolved_kwargs, mode, normalized_legacy_trigger def _gemma4_nonthinking_generation_suffix( @@ -1251,7 +1331,19 @@ def _resolve_assistant_prompt_prefixes( probe_contexts = _assistant_prompt_probe_contexts(messages) if not probe_contexts: return () - cache_key = f"{_serialize_tools_for_cache(tools)}::{','.join(probe_contexts)}" + try: + serialized_template_kwargs = json.dumps( + chat_template_kwargs, + sort_keys=True, + separators=(",", ":"), + default=repr, + ) + except TypeError: + serialized_template_kwargs = repr(chat_template_kwargs) + cache_key = ( + f"{_serialize_tools_for_cache(tools)}::{serialized_template_kwargs}::" + f"{','.join(probe_contexts)}" + ) prefixes = cache.get(cache_key) if prefixes is None: prefixes = _infer_assistant_prompt_prefixes(renderer, tools, chat_template_kwargs, probe_contexts) @@ -2134,7 +2226,9 @@ def _render_training_row( assistant_prompt_prefix_cache: dict[str, tuple[str, ...]], strict: bool, ) -> _RenderedRow | None: - _validate_gemma4_thinking_contract(renderer, messages, template_kwargs) + messages, template_kwargs, gemma4_mode, normalized_legacy_trigger = ( + _resolve_gemma4_thinking_contract(renderer, messages, template_kwargs) + ) if teich_masking: text, supervised_spans = _supervised_text_and_spans( renderer, @@ -2162,6 +2256,8 @@ def _render_training_row( supervised_spans=supervised_spans, tokenized=tokenized, token_length=token_length, + gemma4_mode=gemma4_mode, + normalized_legacy_trigger=normalized_legacy_trigger, ) @@ -2198,6 +2294,11 @@ def row_fits_context( messages = _normalize_tool_call_arguments_for_template(normalize_training_messages(messages)) renderer = _resolve_chat_template_renderer(tokenizer, text_tokenizer) template_kwargs = _validate_chat_template_kwargs(chat_template_kwargs) + messages, template_kwargs, _, _ = _resolve_gemma4_thinking_contract( + renderer, + messages, + template_kwargs, + ) text = _render_chat(renderer, messages, tools, template_kwargs) token_length = _tokenized_length(text_tokenizer, text) result = RowContextFit( @@ -2479,6 +2580,10 @@ def render_candidate(candidate_index: int) -> _RenderedRow | None: ) if report is not None: report.record_kept_row(row_info, rendered.token_length) + report.record_gemma4_mode( + rendered.gemma4_mode, + normalized_legacy_trigger=rendered.normalized_legacy_trigger, + ) return output_batch formatted_data = dataset.map( diff --git a/tests/test_audit.py b/tests/test_audit.py index 40661ab..7398a4a 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -132,6 +132,10 @@ def test_gemma4_example_uses_live_remote_template_and_safe_masks(): assert 'os.environ.setdefault("UNSLOTH_RETURN_LOGITS", "1")' in source assert 'MODEL_REVISION = os.environ.get("MODEL_REVISION", "main")' in source assert 'CHAT_TEMPLATE_PATH = os.environ.get("CHAT_TEMPLATE_PATH")' in source + assert 'os.environ.get("GEMMA4_THINKING_MODE")' in source + assert 'os.environ.get("GEMMA4_ENABLE_THINKING")' in source + assert 'GEMMA4_THINKING_MODE not in {"auto", "thinking", "nonthinking"}' in source + assert "chat_template_kwargs=CHAT_TEMPLATE_KWARGS" in source assert 'or "gemma-template.jinja"' not in source assert "token=HF_TOKEN or None" in source assert 'oversized_policy="trim_followups"' in source diff --git a/tests/test_extract_anonymize_cli.py b/tests/test_extract_anonymize_cli.py index a88ca62..7f64a10 100644 --- a/tests/test_extract_anonymize_cli.py +++ b/tests/test_extract_anonymize_cli.py @@ -1833,7 +1833,17 @@ def test_anonymize_parallel_worker_count_caps_windows(monkeypatch): assert anonymize_module._process_worker_count(100) == 61 -def test_anonymize_path_progress_reports_parallel_completions_and_sorts_final_report(tmp_path: Path): +def test_anonymize_large_batch_uses_process_workers(tmp_path: Path): + sources = [] + for index in range(8): + source = tmp_path / f"large-{index}.jsonl" + source.write_bytes(b"x" * (128 * 1024)) + sources.append(source) + + assert anonymize_module._use_process_workers(sources, workers=8) is True + + +def test_anonymize_path_small_batch_avoids_process_startup_and_reports_progress(tmp_path: Path): input_dir = tmp_path / "input" input_dir.mkdir() for index in reversed(range(9)): @@ -1843,11 +1853,16 @@ def test_anonymize_path_progress_reports_parallel_completions_and_sorts_final_re ) updates = [] - report = anonymize_module.anonymize_path( - input_dir, - tmp_path / "output", - progress=lambda file_report, done, total: updates.append((file_report, done, total)), - ) + with patch.object( + anonymize_module, + "ProcessPoolExecutor", + side_effect=AssertionError("tiny files should stay in-process"), + ): + report = anonymize_module.anonymize_path( + input_dir, + tmp_path / "output", + progress=lambda file_report, done, total: updates.append((file_report, done, total)), + ) assert {update[0].path.name for update in updates} == {f"trace-{index}.jsonl" for index in range(9)} assert [update[1] for update in updates] == list(range(1, 10)) diff --git a/tests/test_formatter.py b/tests/test_formatter.py index e09fa4b..e2d634a 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -396,6 +396,7 @@ def apply_chat_template( add_generation_prompt=False, tools=None, enable_thinking=True, + preserve_thinking=False, **kwargs, ): self.render_count += 1 @@ -474,8 +475,8 @@ def test_gemma4_rejects_reasoning_when_thinking_is_disabled(): ) -@pytest.mark.parametrize("enable_thinking", [False, True]) -def test_gemma4_rejects_manual_think_trigger_in_source_data(enable_thinking: bool): +@pytest.mark.parametrize("enable_thinking", [None, True]) +def test_gemma4_normalizes_leading_legacy_think_trigger(enable_thinking: bool | None): tokenizer = GemmaLikeOffsetTokenizer() tokenizer.name_or_path = "google/gemma-4-E4B-it" dataset = Dataset.from_list( @@ -491,19 +492,126 @@ def test_gemma4_rejects_manual_think_trigger_in_source_data(enable_thinking: boo ] ) - with pytest.raises(ValueError, match="Do not embed the trigger manually"): + template_kwargs = {} if enable_thinking is None else {"enable_thinking": enable_thinking} + prepared, report = prepare_data( + dataset, + tokenizer, + chat_template_kwargs=template_kwargs, + strict=True, + return_report=True, + verbose=False, + ) + + assert prepared.num_rows == 1 + assert "<|think|>" not in prepared[0]["text"] + assert report.gemma4_modes == {"thinking": 1} + assert report.gemma4_legacy_triggers_normalized == 1 + + +def test_gemma4_rejects_legacy_trigger_when_explicitly_disabled(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-E4B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "system", "content": "<|think|>"}, + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ], + "tools": [], + } + ] + ) + + with pytest.raises(ValueError, match="enable_thinking=False conflicts"): prepare_data( dataset, tokenizer, - chat_template_kwargs={ - "enable_thinking": enable_thinking, - "preserve_thinking": enable_thinking, - }, + chat_template_kwargs={"enable_thinking": False}, strict=True, verbose=False, ) +def test_gemma4_auto_mode_handles_mixed_thinking_and_direct_rows(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-26B-A4B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "reason"}, + {"role": "assistant", "content": "answer", "reasoning_content": "analysis"}, + ], + "tools": [], + }, + { + "messages": [ + {"role": "user", "content": "direct"}, + {"role": "assistant", "content": "answer"}, + ], + "tools": [], + }, + ] + ) + + prepared, report = prepare_data( + dataset, + tokenizer, + strict=True, + return_report=True, + verbose=False, + ) + + assert "<|channel>thought\nanalysis\n" in prepared[0]["text"] + assert "<|channel>thought\nanswer" in prepared[1]["text"] + assert report.gemma4_modes == {"thinking": 1, "nonthinking": 1} + + training_data = prepare_and_mask_for_test(dataset, tokenizer, strict=True) + thinking_target = tokenizer.decode( + [token for token in training_data[0]["labels"] if token != -100] + ) + direct_target = tokenizer.decode( + [token for token in training_data[1]["labels"] if token != -100] + ) + assert "analysis" in thinking_target + assert thinking_target.endswith("") + assert "<|channel>thought" not in direct_target + assert "answer" in direct_target + assert direct_target.endswith("") + + +def test_gemma4_auto_mode_preserves_historical_reasoning(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-31B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "one", "reasoning_content": "reason one"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "two", "reasoning_content": "reason two"}, + ], + "tools": [], + } + ] + ) + + prepared, report = prepare_data( + dataset, + tokenizer, + strict=True, + return_report=True, + verbose=False, + ) + + assert "reason one" in prepared[0]["text"] + assert "reason two" in prepared[0]["text"] + assert report.gemma4_modes == {"thinking": 1} + + def test_gemma4_rejects_dropped_historical_reasoning(): tokenizer = GemmaLikeOffsetTokenizer() tokenizer.name_or_path = "google/gemma-4-31B-it" diff --git a/tests/test_tokenizer_smoke.py b/tests/test_tokenizer_smoke.py index ccb7e9a..bac0554 100644 --- a/tests/test_tokenizer_smoke.py +++ b/tests/test_tokenizer_smoke.py @@ -13,17 +13,17 @@ pytest.param("unsloth/Qwen3.5-0.8B", {"enable_thinking": True}, id="unsloth-qwen3.5"), pytest.param( "google/gemma-4-E4B-it", - {"enable_thinking": True, "preserve_thinking": True}, + {}, id="gemma-4-e4b-it", ), pytest.param( "google/gemma-4-26B-A4B-it", - {"enable_thinking": True, "preserve_thinking": True}, + {}, id="gemma-4-26b-a4b-it", ), pytest.param( "google/gemma-4-31B-it", - {"enable_thinking": True, "preserve_thinking": True}, + {}, id="gemma-4-31b-it", ), ] @@ -116,6 +116,7 @@ def test_real_tokenizer_prepare_and_mask_tool_dataset(model_id: str, chat_templa assert prepared.column_names == ["text", "teich_supervised_spans", "input_ids", "attention_mask"] assert trainer.train_dataset.column_names == ["input_ids", "labels"] assert supervised_ids + assert "I should inspect the workspace." in supervised_text assert "bash" in supervised_text assert "ls" in supervised_text assert "Found project files." in supervised_text @@ -130,12 +131,13 @@ def test_real_tokenizer_prepare_and_mask_tool_dataset(model_id: str, chat_templa source_row = _tool_call_dataset()[0] generation_messages = source_row["messages"][:-1] + generation_kwargs = {"enable_thinking": True, "preserve_thinking": True} generation_prompt = tokenizer.apply_chat_template( generation_messages, tools=source_row["tools"], tokenize=False, add_generation_prompt=True, - **chat_template_kwargs, + **generation_kwargs, ) assert generation_prompt.endswith("<|channel>thought\n") @@ -168,7 +170,6 @@ def test_real_gemma4_nonthinking_training_matches_generation_prompt( tokenizer, tokenize=True, strict=True, - chat_template_kwargs={"enable_thinking": False}, verbose=False, ) empty_thought = "<|channel>thought\n" @@ -182,3 +183,19 @@ def test_real_gemma4_nonthinking_training_matches_generation_prompt( assert (empty_thought in generation_prompt) is expects_empty_thought assert (empty_thought in prepared[0]["text"]) is expects_empty_thought assert "<|think|>" not in prepared[0]["text"] + + trainer = SimpleNamespace( + train_dataset=prepared, + eval_dataset=None, + processing_class=tokenizer, + args=SimpleNamespace(dataset_text_field="text", packing=False, max_length=4096), + ) + trainer = mask_data(trainer, tokenizer=tokenizer, audit=True, verbose=False) + row = trainer.train_dataset[0] + supervised_text = tokenizer.decode( + [token for token in row["labels"] if token != -100], + skip_special_tokens=False, + ) + assert empty_thought not in supervised_text + assert "Direct answer." in supervised_text + assert supervised_text.endswith("") From 270fb6d1d313e41f28a45c028d2f8a17641547ab Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Tue, 25 Aug 2026 00:49:10 -0400 Subject: [PATCH 3/6] Support template-native reasoning training policies --- README.md | 2 + docs/prepare-data.md | 8 +- docs/python-api.md | 1 + docs/training.md | 65 +++++++++++ src/teich/formatter.py | 51 +++++++++ src/teich/prepare.py | 23 +++- tests/test_formatter.py | 65 +++++++++++ tests/test_prepare.py | 82 ++++++++++++++ tests/test_tokenizer_smoke.py | 206 ++++++++++++++++++++++++++++++++++ 9 files changed, 500 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 6b7f5c7..041b9a0 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,8 @@ train_dataset = prepare_data( With a live Gemma 4 tokenizer, omit `chat_template_kwargs` to let Teich choose thinking or non-thinking independently for each row. See [Training](docs/training.md#live-gemma-4-models) for auto-mode rules and the E4B, 26B-A4B, and 31B template contract. +Qwen 3.8 keeps its own template-native defaults, including historical reasoning +preservation and `reasoning_effort`; see [Live Qwen 3.8 Models](docs/training.md#live-qwen-38-models). Then create your trainer and call `mask_data()`: diff --git a/docs/prepare-data.md b/docs/prepare-data.md index 0053f37..ea10515 100644 --- a/docs/prepare-data.md +++ b/docs/prepare-data.md @@ -51,7 +51,9 @@ print(prep_report.oversized_rows[:3]) `PrepareReport` includes dropped rows, oversized rows, trimmed rows, token lengths, max token length, kept-row ids, and returned row count. With a live Gemma 4 template it also reports per-row mode counts in `gemma4_modes` and migrated -leading `<|think|>` markers in `gemma4_legacy_triggers_normalized`. +leading `<|think|>` markers in `gemma4_legacy_triggers_normalized`. When +`reasoning_policy="strip"` is used, it reports affected rows and messages in +`reasoning_stripped_rows` and `reasoning_stripped_messages`. Gemma 4 defaults to per-row auto mode when `enable_thinking` is omitted: reasoning-bearing rows enable thinking and history preservation, while direct @@ -89,6 +91,7 @@ train_dataset = prepare_data( "instruct-chat": { "source": "TeichAI/polaris-alpha-1000x", "percentage": 20, + "reasoning_policy": "strip", "chat_template_kwargs": {"enable_thinking": False, "preserve_thinking": False}, }, }, @@ -104,7 +107,8 @@ train_dataset = prepare_data( If one source cannot fill its share after filtering or context-window drops, Teich scales the total row count down instead of silently changing the realized mix. -Global `chat_template_kwargs` are the default for every source. A source-level `chat_template_kwargs` mapping overrides those keys for that dataset only. +Global `chat_template_kwargs` and `reasoning_policy` values are the defaults for +every source. Source-level values override them for that dataset only. You can also pass a simple list of sources: diff --git a/docs/python-api.md b/docs/python-api.md index b84a7e5..74490a5 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -65,6 +65,7 @@ Useful options: - `teich_masking` - `tokenize` - `chat_template_kwargs` +- `reasoning_policy` (`"keep"` or `"strip"`) See [Preparing Data](prepare-data.md). diff --git a/docs/training.md b/docs/training.md index c63b0e4..336dc11 100644 --- a/docs/training.md +++ b/docs/training.md @@ -121,6 +121,65 @@ environment and check it with `python -m pip check` before a long run. Teich's core environment does not pin the CUDA, PyTorch, Unsloth, and TRL stack because those versions depend on the host GPU and CUDA runtime. +## Live Qwen 3.8 Models + +Qwen 3.8 has a different native contract and Teich does not apply Gemma's auto +mode rules to it. The live `Qwen/Qwen3.8-27B` template defaults to thinking, +defaults `reasoning_effort` to `xhigh`, and preserves historical +`reasoning_content` unless `preserve_thinking=False` is supplied. To train its +native reasoning behavior, retain those defaults or set only the desired +effort: + +```python +train_dataset = prepare_data( + "username/qwen38-reasoning-traces", + tokenizer, + chat_template_kwargs={"reasoning_effort": "medium"}, + tokenize=True, + strict=True, +) +``` + +Supported live efforts are `low`, `medium`, and `xhigh`. Continue to use +`train_on_reasoning=True` in `mask_data()` when those reasoning tokens should +receive loss. + +For direct instruction tuning from a dataset that still contains reasoning, +remove the reasoning before rendering and explicitly select Qwen's +non-thinking template mode: + +```python +train_dataset, prep_report = prepare_data( + "username/mixed-source-traces", + tokenizer, + reasoning_policy="strip", + chat_template_kwargs={ + "enable_thinking": False, + "preserve_thinking": False, + }, + return_report=True, + tokenize=True, + strict=True, +) +``` + +Qwen 3.8's non-thinking prompt contains an empty `...` primer. +Teich keeps that inference-alignment prefix in the rendered text but masks it +from loss. The final answer and closing `<|im_end|>` remain supervised. + +## Source Reasoning Policy + +`reasoning_policy="keep"` is the default and leaves structured assistant +reasoning for the loaded chat template to handle according to its own model +contract. `reasoning_policy="strip"` removes normalized `reasoning_content` +before rendering. This is deliberately independent of +`mask_data(train_on_reasoning=False)`: masking keeps gold reasoning in the +causal context, whereas stripping creates a true instruction-only example. + +The policy can be set per source in a mixed dataset. `PrepareReport` records +the affected row and message counts in `reasoning_stripped_rows` and +`reasoning_stripped_messages`. + ## What `mask_data()` Does Before `mask_data()`, the trainer dataset usually contains: @@ -155,6 +214,12 @@ each dataset-map batch together instead of issuing one tokenizer call per row. For Qwen-style templates, the initial `` tag is intentionally included in supervision. +For Gemma 4, Teich supervises exactly one closing `` for a completed +model turn whenever reasoning, final-answer, or tool-call supervision is +enabled for that turn. It does not add a second terminator inside a continuing +tool-call chain. The terminator remains a target even when the final answer is +masked, so reasoning-only and tool-only fine-tunes still learn to stop. + ## Masking Policy `mask_data()` trains on these by default: diff --git a/src/teich/formatter.py b/src/teich/formatter.py index 4f6af58..14ff223 100644 --- a/src/teich/formatter.py +++ b/src/teich/formatter.py @@ -75,6 +75,7 @@ } _OVERSIZED_POLICIES = {"drop", "trim_followups", "error"} _OVERSIZED_POLICY_KEEP = "keep" +_REASONING_POLICIES = {"keep", "strip"} _GEMMA4_PREFIX_CACHE: weakref.WeakKeyDictionary[Any, dict[str, str]] = weakref.WeakKeyDictionary() @@ -100,6 +101,8 @@ class PrepareReport: trimmed_rows: list[dict[str, Any]] = field(default_factory=list) gemma4_modes: dict[str, int] = field(default_factory=dict) gemma4_legacy_triggers_normalized: int = 0 + reasoning_stripped_rows: int = 0 + reasoning_stripped_messages: int = 0 def record_token_length(self, row_info: dict[str, Any], token_length: int) -> None: entry = {**row_info, "token_length": token_length} @@ -162,6 +165,12 @@ def record_gemma4_mode(self, mode: str | None, *, normalized_legacy_trigger: boo if normalized_legacy_trigger: self.gemma4_legacy_triggers_normalized += 1 + def record_stripped_reasoning(self, stripped_messages: int) -> None: + if stripped_messages <= 0: + return + self.reasoning_stripped_rows += 1 + self.reasoning_stripped_messages += stripped_messages + def to_dict(self) -> dict[str, Any]: return { "total_rows": self.total_rows, @@ -176,6 +185,8 @@ def to_dict(self) -> dict[str, Any]: "trimmed_rows": self.trimmed_rows, "gemma4_modes": self.gemma4_modes, "gemma4_legacy_triggers_normalized": self.gemma4_legacy_triggers_normalized, + "reasoning_stripped_rows": self.reasoning_stripped_rows, + "reasoning_stripped_messages": self.reasoning_stripped_messages, } @@ -225,6 +236,35 @@ def _validate_chat_template_kwargs(chat_template_kwargs: dict[str, Any] | None) return kwargs +def _validate_reasoning_policy(reasoning_policy: str) -> str: + if not isinstance(reasoning_policy, str) or reasoning_policy not in _REASONING_POLICIES: + choices = ", ".join(sorted(_REASONING_POLICIES)) + raise ValueError(f"reasoning_policy must be one of: {choices}.") + return reasoning_policy + + +def _apply_reasoning_policy( + messages: list[dict[str, Any]], + reasoning_policy: str, +) -> tuple[list[dict[str, Any]], int]: + if reasoning_policy == "keep": + return messages, 0 + stripped_messages = 0 + resolved_messages = messages + for index, message in enumerate(messages): + if message.get("role") not in {"assistant", "model"} or not _message_reasoning(message): + continue + if resolved_messages is messages: + resolved_messages = list(messages) + updated = dict(message) + updated.pop("reasoning_content", None) + updated.pop("thinking", None) + updated.pop("reasoning", None) + resolved_messages[index] = updated + stripped_messages += 1 + return resolved_messages, stripped_messages + + def _renderer_name(renderer: Any) -> str: for candidate in (renderer, getattr(renderer, "tokenizer", None)): name = getattr(candidate, "name_or_path", None) @@ -2270,6 +2310,7 @@ def row_fits_context( messages_column: str = "messages", tools_column: str = "tools", text_column: str = "text", + reasoning_policy: str = "keep", return_details: bool = False, ) -> bool | RowContextFit: if not isinstance(max_length, int) or max_length <= 0: @@ -2292,6 +2333,7 @@ def row_fits_context( if not isinstance(tools, list): raise TypeError(f"Row has a non-list '{tools_column}' column.") messages = _normalize_tool_call_arguments_for_template(normalize_training_messages(messages)) + messages, _ = _apply_reasoning_policy(messages, _validate_reasoning_policy(reasoning_policy)) renderer = _resolve_chat_template_renderer(tokenizer, text_tokenizer) template_kwargs = _validate_chat_template_kwargs(chat_template_kwargs) messages, template_kwargs, _, _ = _resolve_gemma4_thinking_contract( @@ -2318,6 +2360,7 @@ def format_data( tools_column: str = "tools", text_column: str = "text", chat_template_kwargs: dict[str, Any] | None = None, + reasoning_policy: str = "keep", train_on_reasoning: bool | None = None, teich_masking: bool = True, max_length: int | None = None, @@ -2349,6 +2392,7 @@ def format_data( tools_column=tools_column, text_column=text_column, chat_template_kwargs=chat_template_kwargs, + reasoning_policy=reasoning_policy, train_on_reasoning=train_on_reasoning, teich_masking=teich_masking, max_length=max_length, @@ -2372,6 +2416,7 @@ def format_data( raise TypeError("prepare_data expects a Dataset or a sequence of Dataset objects.") template_kwargs = _validate_chat_template_kwargs(chat_template_kwargs) + effective_reasoning_policy = _validate_reasoning_policy(reasoning_policy) text_tokenizer = _resolve_text_tokenizer(tokenizer) renderer = _resolve_chat_template_renderer(tokenizer, text_tokenizer) assistant_prompt_prefix_cache: dict[str, tuple[str, ...]] = {} @@ -2430,6 +2475,12 @@ def _map_batch(batch: dict[str, list[Any]], indices: list[int]) -> dict[str, lis report.record_dropped_row(row_info, "empty_messages") continue messages = _normalize_tool_call_arguments_for_template(messages) + messages, stripped_reasoning_messages = _apply_reasoning_policy( + messages, + effective_reasoning_policy, + ) + if report is not None: + report.record_stripped_reasoning(stripped_reasoning_messages) tools = tools_batch[index] or [] if not isinstance(tools, list): raise TypeError(f"Row is missing a list-valued '{tools_column}' column") diff --git a/src/teich/prepare.py b/src/teich/prepare.py index 66a4986..3e23f8d 100644 --- a/src/teich/prepare.py +++ b/src/teich/prepare.py @@ -23,6 +23,7 @@ class _SourceMixEntry: percentage: float | None has_explicit_mix_value: bool chat_template_kwargs: dict[str, Any] | None + reasoning_policy: str | None @dataclass(slots=True) @@ -33,6 +34,7 @@ class _ResolvedSourceMix: names: list[str] rigid_percentages: bool chat_template_kwargs: list[dict[str, Any] | None] + reasoning_policies: list[str | None] @dataclass(slots=True) @@ -57,6 +59,7 @@ def prepare_data( tools_column: str = "tools", text_column: str = "text", chat_template_kwargs: dict[str, Any] | None = None, + reasoning_policy: str = "keep", train_on_reasoning: bool | None = None, teich_masking: bool = True, max_length: int | None = None, @@ -94,6 +97,7 @@ def prepare_data( tools_column=tools_column, text_column=text_column, chat_template_kwargs=_merge_chat_template_kwargs(chat_template_kwargs, source_chat_template_kwargs), + reasoning_policy=source_reasoning_policy or reasoning_policy, train_on_reasoning=train_on_reasoning, teich_masking=teich_masking, max_length=max_length, @@ -108,9 +112,10 @@ def prepare_data( strict=strict, verbose=verbose, ) - for source_dataset, source_chat_template_kwargs, source_name in zip( + for source_dataset, source_chat_template_kwargs, source_reasoning_policy, source_name in zip( dataset.datasets, dataset.chat_template_kwargs, + dataset.reasoning_policies, dataset.names, strict=True, ) @@ -134,6 +139,7 @@ def prepare_data( tools_column=tools_column, text_column=text_column, chat_template_kwargs=chat_template_kwargs, + reasoning_policy=reasoning_policy, train_on_reasoning=train_on_reasoning, teich_masking=teich_masking, max_length=max_length, @@ -171,6 +177,7 @@ def prepare_data( tools_column=tools_column, text_column=text_column, chat_template_kwargs=chat_template_kwargs, + reasoning_policy=reasoning_policy, train_on_reasoning=train_on_reasoning, teich_masking=teich_masking, max_length=max_length, @@ -306,6 +313,10 @@ def _source_mix_entry_from_value(value: Any, *, default_name: str) -> _SourceMix value.get("chat_template_kwargs"), f"{name_value}.chat_template_kwargs", ), + reasoning_policy=_optional_reasoning_policy( + value.get("reasoning_policy"), + f"{name_value}.reasoning_policy", + ), ) raise TypeError("A source mix entry mapping must include a 'source', 'dataset', or 'path' key.") return _SourceMixEntry( @@ -315,6 +326,7 @@ def _source_mix_entry_from_value(value: Any, *, default_name: str) -> _SourceMix percentage=None, has_explicit_mix_value=False, chat_template_kwargs=None, + reasoning_policy=None, ) @@ -382,6 +394,14 @@ def _optional_chat_template_kwargs(value: Any, name: str) -> dict[str, Any] | No return dict(value) +def _optional_reasoning_policy(value: Any, name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str) or value not in {"keep", "strip"}: + raise ValueError(f"{name} must be either 'keep' or 'strip'.") + return value + + def _merge_chat_template_kwargs( global_kwargs: dict[str, Any] | None, source_kwargs: dict[str, Any] | None, @@ -439,6 +459,7 @@ def _resolve_source_mix( names=[entry.name for entry in entries], rigid_percentages=any(entry.has_explicit_mix_value for entry in entries), chat_template_kwargs=[entry.chat_template_kwargs for entry in entries], + reasoning_policies=[entry.reasoning_policy for entry in entries], ) diff --git a/tests/test_formatter.py b/tests/test_formatter.py index e2d634a..abcf1ee 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -26,6 +26,7 @@ def prepare_and_mask_for_test( messages_column="messages", tools_column="tools", chat_template_kwargs=None, + reasoning_policy="keep", train_on_reasoning=True, train_on_final_answers=True, train_on_tools=True, @@ -41,6 +42,7 @@ def prepare_and_mask_for_test( messages_column=messages_column, tools_column=tools_column, chat_template_kwargs=chat_template_kwargs, + reasoning_policy=reasoning_policy, train_on_reasoning=train_on_reasoning, max_length=max_length, drop_oversized_examples=drop_oversized_examples, @@ -582,6 +584,69 @@ def test_gemma4_auto_mode_handles_mixed_thinking_and_direct_rows(): assert direct_target.endswith("") +def test_reasoning_policy_strip_makes_reasoning_dataset_direct_for_gemma4(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-26B-A4B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer", "reasoning_content": "private analysis"}, + ], + "tools": [], + } + ] + ) + + prepared, report = prepare_data( + dataset, + tokenizer, + reasoning_policy="strip", + strict=True, + return_report=True, + verbose=False, + ) + + assert "private analysis" not in prepared[0]["text"] + assert "<|channel>thought\nanswer" in prepared[0]["text"] + assert report.gemma4_modes == {"nonthinking": 1} + assert report.reasoning_stripped_rows == 1 + assert report.reasoning_stripped_messages == 1 + + training_data = prepare_and_mask_for_test( + dataset, + tokenizer, + reasoning_policy="strip", + strict=True, + ) + supervised_text = tokenizer.decode( + [token for token in training_data[0]["labels"] if token != -100] + ) + assert "private analysis" not in supervised_text + assert supervised_text == "answer" + + +def test_prepare_data_rejects_unknown_reasoning_policy(): + with pytest.raises(ValueError, match="reasoning_policy must be one of"): + prepare_data( + Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ], + "tools": [], + } + ] + ), + GemmaLikeOffsetTokenizer(), + reasoning_policy="guess", + verbose=False, + ) + + def test_gemma4_auto_mode_preserves_historical_reasoning(): tokenizer = GemmaLikeOffsetTokenizer() tokenizer.name_or_path = "google/gemma-4-31B-it" diff --git a/tests/test_prepare.py b/tests/test_prepare.py index dba383a..8021b6a 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -67,6 +67,22 @@ def apply_chat_template(self, messages, *, tokenize=False, add_generation_prompt return rendered +class ReasoningAwareTokenizer(TinyChatTokenizer): + def apply_chat_template(self, messages, *, tokenize=False, add_generation_prompt=False, tools=None, **kwargs): + parts = [] + for message in messages: + role = message["role"] + reasoning = message.get("reasoning_content") + reasoning_text = f"{reasoning}" if reasoning else "" + parts.append(f"<{role}>{reasoning_text}{message.get('content', '')}") + rendered = "".join(parts) + if add_generation_prompt: + rendered += "" + if tokenize: + return self(rendered) + return rendered + + def _dataset() -> Dataset: return Dataset.from_list( [ @@ -392,6 +408,72 @@ def test_prepare_data_source_mix_rejects_invalid_dataset_level_chat_template_kwa ) +def test_prepare_data_source_mix_supports_dataset_level_reasoning_policy(): + reasoning_dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "reasoning prompt"}, + { + "role": "assistant", + "content": "reasoning answer", + "reasoning_content": "keep this analysis", + }, + ], + "tools": [], + } + ] + ) + instruct_dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "instruct prompt"}, + { + "role": "assistant", + "content": "instruct answer", + "reasoning_content": "strip this analysis", + }, + ], + "tools": [], + } + ] + ) + + prepared, report = prepare_data( + { + "reasoning": {"source": reasoning_dataset, "reasoning_policy": "keep"}, + "instruct": {"source": instruct_dataset, "reasoning_policy": "strip"}, + }, + ReasoningAwareTokenizer(), + return_report=True, + verbose=False, + ) + + texts = [prepared[index]["text"] for index in range(prepared.num_rows)] + reasoning_text = next(text for text in texts if "reasoning answer" in text) + instruct_text = next(text for text in texts if "instruct answer" in text) + assert "keep this analysis" in reasoning_text + assert "strip this analysis" not in instruct_text + assert report.reasoning_stripped_rows == 1 + assert report.reasoning_stripped_messages == 1 + + +@pytest.mark.parametrize("reasoning_policy", [True, "drop", ["strip"]]) +def test_prepare_data_source_mix_rejects_invalid_reasoning_policy(reasoning_policy): + with pytest.raises(ValueError, match="reasoning_policy must be either"): + prepare_data( + { + "bad": { + "source": _dataset(), + "reasoning_policy": reasoning_policy, + } + }, + TinyChatTokenizer(), + verbose=False, + ) + + def test_prepare_data_source_mix_applies_each_source_tools_snapshot_independently(tmp_path: Path): tokenizer = TinyChatTokenizer() alpha_source = tmp_path / "alpha" diff --git a/tests/test_tokenizer_smoke.py b/tests/test_tokenizer_smoke.py index bac0554..25c3118 100644 --- a/tests/test_tokenizer_smoke.py +++ b/tests/test_tokenizer_smoke.py @@ -11,6 +11,7 @@ TOKENIZER_SMOKE_MODELS = [ pytest.param("unsloth/Qwen3.5-0.8B", {"enable_thinking": True}, id="unsloth-qwen3.5"), + pytest.param("Qwen/Qwen3.8-27B", {}, id="qwen3.8-27b"), pytest.param( "google/gemma-4-E4B-it", {}, @@ -123,11 +124,16 @@ def test_real_tokenizer_prepare_and_mask_tool_dataset(model_id: str, chat_templa assert "SECRET_TOOL_OUTPUT" not in supervised_text assert "SECRET_TOOL_OUTPUT" in masked_text + if model_id == "Qwen/Qwen3.8-27B": + assert "Reasoning effort is set to xhigh." in prepared[0]["text"] + assert prepared[0]["text"].count("I should inspect the workspace.") == 1 + if model_id.startswith("google/gemma-4-"): assert "<|think|>" in prepared[0]["text"] eot_token_id = tokenizer.convert_tokens_to_ids("") assert eot_token_id in supervised_ids assert supervised_text.endswith("") + assert supervised_text.count("") == 1 source_row = _tool_call_dataset()[0] generation_messages = source_row["messages"][:-1] @@ -199,3 +205,203 @@ def test_real_gemma4_nonthinking_training_matches_generation_prompt( assert empty_thought not in supervised_text assert "Direct answer." in supervised_text assert supervised_text.endswith("") + assert supervised_text.count("") == 1 + + +@pytest.mark.integration +@pytest.mark.tokenizer_smoke +@pytest.mark.parametrize( + "model_id", + [ + pytest.param("google/gemma-4-E4B-it", id="gemma-4-e4b-turn-end"), + pytest.param("google/gemma-4-26B-A4B-it", id="gemma-4-26b-a4b-turn-end"), + pytest.param("google/gemma-4-31B-it", id="gemma-4-31b-turn-end"), + ], +) +def test_real_gemma4_supervises_exactly_one_turn_end_for_enabled_targets(model_id: str): + if not _tokenizer_smokes_enabled(): + pytest.skip("Set TEICH_RUN_TOKENIZER_SMOKES=1 to run real Hugging Face tokenizer smokes.") + transformers = pytest.importorskip("transformers") + tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + simple = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "Solve this."}, + { + "role": "assistant", + "content": "Final answer.", + "reasoning_content": "Careful reasoning.", + }, + ], + "tools": [], + } + ] + ) + prepared = prepare_data(simple, tokenizer, tokenize=True, strict=True, verbose=False) + turn_end_token_id = tokenizer.convert_tokens_to_ids("") + + for train_on_reasoning, train_on_final_answers in [(True, False), (False, True)]: + trainer = SimpleNamespace( + train_dataset=prepared, + eval_dataset=None, + processing_class=tokenizer, + args=SimpleNamespace(dataset_text_field="text", packing=False, max_length=4096), + ) + trainer = mask_data( + trainer, + tokenizer=tokenizer, + train_on_reasoning=train_on_reasoning, + train_on_final_answers=train_on_final_answers, + train_on_tools=False, + audit=True, + verbose=False, + ) + supervised_text = tokenizer.decode( + [token for token in trainer.train_dataset[0]["labels"] if token != -100], + skip_special_tokens=False, + ) + supervised_ids = [token for token in trainer.train_dataset[0]["labels"] if token != -100] + assert supervised_ids.count(turn_end_token_id) == 1 + assert supervised_text.count("") == 1 + assert supervised_text.endswith("") + assert ("Careful reasoning." in supervised_text) is train_on_reasoning + assert ("Final answer." in supervised_text) is train_on_final_answers + + tool_prepared = prepare_data( + _tool_call_dataset(), + tokenizer, + tokenize=True, + strict=True, + max_length=4096, + verbose=False, + ) + tool_trainer = SimpleNamespace( + train_dataset=tool_prepared, + eval_dataset=None, + processing_class=tokenizer, + args=SimpleNamespace(dataset_text_field="text", packing=False, max_length=4096), + ) + tool_trainer = mask_data( + tool_trainer, + tokenizer=tokenizer, + train_on_reasoning=False, + train_on_final_answers=False, + train_on_tools=True, + audit=True, + verbose=False, + ) + tool_supervised_text = tokenizer.decode( + [token for token in tool_trainer.train_dataset[0]["labels"] if token != -100], + skip_special_tokens=False, + ) + tool_supervised_ids = [token for token in tool_trainer.train_dataset[0]["labels"] if token != -100] + assert "bash" in tool_supervised_text + assert "ls" in tool_supervised_text + assert "SECRET_TOOL_OUTPUT" not in tool_supervised_text + assert "Found project files." not in tool_supervised_text + assert tool_supervised_ids.count(turn_end_token_id) == 1 + assert tool_supervised_text.count("") == 1 + assert tool_supervised_text.endswith("") + + +@pytest.mark.integration +@pytest.mark.tokenizer_smoke +def test_real_qwen38_preserves_reasoning_history_and_honors_effort(): + if not _tokenizer_smokes_enabled(): + pytest.skip("Set TEICH_RUN_TOKENIZER_SMOKES=1 to run real Hugging Face tokenizer smokes.") + transformers = pytest.importorskip("transformers") + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen3.8-27B", trust_remote_code=True) + messages = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "one", "reasoning_content": "reason one"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "two", "reasoning_content": "reason two"}, + ] + + prepared = prepare_data( + Dataset.from_list([{"messages": messages, "tools": []}]), + tokenizer, + tokenize=True, + strict=True, + chat_template_kwargs={"reasoning_effort": "low"}, + verbose=False, + ) + text = prepared[0]["text"] + assert "Reasoning effort is set to low." in text + assert text.count("reason one") == 1 + assert text.count("reason two") == 1 + + trainer = SimpleNamespace( + train_dataset=prepared, + eval_dataset=None, + processing_class=tokenizer, + args=SimpleNamespace(dataset_text_field="text", packing=False, max_length=4096), + ) + trainer = mask_data(trainer, tokenizer=tokenizer, train_on_reasoning=True, audit=True, verbose=False) + supervised_text = tokenizer.decode( + [token for token in trainer.train_dataset[0]["labels"] if token != -100], + skip_special_tokens=False, + ) + assert "reason one" in supervised_text + assert "reason two" in supervised_text + assert "one" in supervised_text + assert "two" in supervised_text + + generation_prompt = tokenizer.apply_chat_template( + messages[:-1], + tokenize=False, + add_generation_prompt=True, + reasoning_effort="low", + ) + assert "reason one" in generation_prompt + assert generation_prompt.endswith("\n") + + +@pytest.mark.integration +@pytest.mark.tokenizer_smoke +def test_real_qwen38_can_strip_reasoning_for_direct_instruct_training(): + if not _tokenizer_smokes_enabled(): + pytest.skip("Set TEICH_RUN_TOKENIZER_SMOKES=1 to run real Hugging Face tokenizer smokes.") + transformers = pytest.importorskip("transformers") + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen3.8-27B", trust_remote_code=True) + messages = [ + {"role": "user", "content": "Give a direct answer."}, + { + "role": "assistant", + "content": "Direct answer.", + "reasoning_content": "reasoning that should be removed", + }, + ] + + prepared, report = prepare_data( + Dataset.from_list([{"messages": messages, "tools": []}]), + tokenizer, + tokenize=True, + strict=True, + reasoning_policy="strip", + chat_template_kwargs={"enable_thinking": False, "preserve_thinking": False}, + return_report=True, + verbose=False, + ) + empty_thought = "\n\n" + assert "reasoning that should be removed" not in prepared[0]["text"] + assert empty_thought in prepared[0]["text"] + assert report.reasoning_stripped_rows == 1 + assert report.reasoning_stripped_messages == 1 + + trainer = SimpleNamespace( + train_dataset=prepared, + eval_dataset=None, + processing_class=tokenizer, + args=SimpleNamespace(dataset_text_field="text", packing=False, max_length=4096), + ) + trainer = mask_data(trainer, tokenizer=tokenizer, train_on_reasoning=True, audit=True, verbose=False) + supervised_text = tokenizer.decode( + [token for token in trainer.train_dataset[0]["labels"] if token != -100], + skip_special_tokens=False, + ) + assert "reasoning that should be removed" not in supervised_text + assert empty_thought not in supervised_text + assert "Direct answer." in supervised_text + assert supervised_text.endswith("<|im_end|>\n") From 726dece6de7bc2c806af2a9b889fa5ec4a435f2a Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Tue, 25 Aug 2026 01:03:14 -0400 Subject: [PATCH 4/6] Update Gemma 4 training example --- docs/training.md | 8 ++++++++ gemma4_example.py | 31 ++++++++++++++++++++++++++++++- tests/test_audit.py | 8 ++++++++ 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/docs/training.md b/docs/training.md index 336dc11..482c66c 100644 --- a/docs/training.md +++ b/docs/training.md @@ -116,6 +116,14 @@ non-`main` revision is required. Set `HF_TOKEN` to an account that has accepted the Gemma repository terms when the checkpoint is not already available through the local Hugging Face login. +The example keeps reasoning in its agent source and strips reasoning from its +direct-chat source. Override those source policies with +`AGENT_REASONING_POLICY` or `CHAT_REASONING_POLICY` when using datasets with a +different contract. It prints the resolved Gemma mode counts, stripped-row +count, and maximum token length before training. Do not add `` to source +messages: Teich derives and supervises the live template's completed-turn +terminator automatically. + Run the example in a dedicated, internally consistent Unsloth training environment and check it with `python -m pip check` before a long run. Teich's core environment does not pin the CUDA, PyTorch, Unsloth, and TRL stack because diff --git a/gemma4_example.py b/gemma4_example.py index 3061da5..97d467f 100644 --- a/gemma4_example.py +++ b/gemma4_example.py @@ -17,6 +17,14 @@ HUB_REPO_ID = os.environ.get("HUB_REPO_ID") or "" HF_TOKEN = os.environ.get("HF_TOKEN", "") CHAT_TEMPLATE_PATH = os.environ.get("CHAT_TEMPLATE_PATH") +AGENT_REASONING_POLICY = os.environ.get("AGENT_REASONING_POLICY", "keep").strip().lower() +CHAT_REASONING_POLICY = os.environ.get("CHAT_REASONING_POLICY", "strip").strip().lower() +for policy_name, policy in { + "AGENT_REASONING_POLICY": AGENT_REASONING_POLICY, + "CHAT_REASONING_POLICY": CHAT_REASONING_POLICY, +}.items(): + if policy not in {"keep", "strip"}: + raise ValueError(f"{policy_name} must be keep or strip") _thinking_mode = os.environ.get("GEMMA4_THINKING_MODE") _legacy_thinking = os.environ.get("GEMMA4_ENABLE_THINKING") if _thinking_mode is None and _legacy_thinking is not None: @@ -66,16 +74,23 @@ random_state = 3407, ) -train_dataset = prepare_data( +train_dataset, prep_report = prepare_data( { "max_examples": 30, "agent": { "source": "armand0e/ag-datagen-v2-test", "percentage": 80, + # Keep structured reasoning. In auto mode, Teich renders these as + # thinking rows and preserves reasoning across multi-turn history. + "reasoning_policy": AGENT_REASONING_POLICY, }, "chat": { "source": "armand0e/DeepSeek-v4-Flash-Chat", "percentage": 20, + # Make this a true direct-instruction source even if an upstream + # row happens to contain reasoning fields. This differs from only + # masking reasoning loss, which would leave it in causal context. + "reasoning_policy": CHAT_REASONING_POLICY, }, }, tokenizer, @@ -89,6 +104,16 @@ oversized_policy="trim_followups", tokenize=True, strict=True, + return_report=True, +) + +print( + "Prepared Gemma 4 modes:", + prep_report.gemma4_modes, + "| stripped reasoning rows:", + prep_report.reasoning_stripped_rows, + "| max tokens:", + prep_report.max_token_length, ) trainer = SFTTrainer( @@ -127,6 +152,10 @@ train_on_tools=True, ) +# Teich keeps exactly one target for each completed Gemma model turn +# that has an enabled reasoning, answer, or tool-call target. Do not append a +# terminator to dataset content manually. + print(trainer.train_dataset.preview()) trainer_stats = trainer.train(resume_from_checkpoint=False) diff --git a/tests/test_audit.py b/tests/test_audit.py index 7398a4a..24bfefe 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -136,6 +136,14 @@ def test_gemma4_example_uses_live_remote_template_and_safe_masks(): assert 'os.environ.get("GEMMA4_ENABLE_THINKING")' in source assert 'GEMMA4_THINKING_MODE not in {"auto", "thinking", "nonthinking"}' in source assert "chat_template_kwargs=CHAT_TEMPLATE_KWARGS" in source + assert '"reasoning_policy": AGENT_REASONING_POLICY' in source + assert '"reasoning_policy": CHAT_REASONING_POLICY' in source + assert 'AGENT_REASONING_POLICY = os.environ.get("AGENT_REASONING_POLICY", "keep")' in source + assert 'CHAT_REASONING_POLICY = os.environ.get("CHAT_REASONING_POLICY", "strip")' in source + assert "train_dataset, prep_report = prepare_data(" in source + assert "return_report=True" in source + assert "prep_report.gemma4_modes" in source + assert "Do not append a" in source assert 'or "gemma-template.jinja"' not in source assert "token=HF_TOKEN or None" in source assert 'oversized_policy="trim_followups"' in source From a42b97b8afa85e88df7e6a334097d7a13c38e301 Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Tue, 25 Aug 2026 01:17:45 -0400 Subject: [PATCH 5/6] Fix connector review findings --- docs/cli.md | 2 +- docs/generation.md | 2 +- src/teich/anonymize.py | 55 ++++++++++++++++++++++------- src/teich/formatter.py | 22 ++++++++++-- src/teich/studio/interactive.py | 5 +++ src/teich/studio/server.py | 47 ++++++++++++++++++------ src/teich/studio/static/app.js | 1 + tests/test_extract_anonymize_cli.py | 44 ++++++++++++++++++++++- tests/test_formatter.py | 42 ++++++++++++++++++++++ tests/test_studio.py | 37 +++++++++++++++++++ 10 files changed, 227 insertions(+), 30 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index b853511..f566a5c 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -98,7 +98,7 @@ teich anonymize output --output output_anonymized teich anonymize data --in-place ``` -Anonymization replaces known credential formats, high-confidence secret assignments, personal email addresses, contextual PII, home-directory usernames, and embedded base64 media with deterministic dummy values. Reserved example-domain addresses, known public bot addresses, provider thinking signatures, placeholders, and common public IDs are preserved to avoid corrupting training data. User, assistant, reasoning, and structured tool content all pass through the same high-confidence privacy scanner. Reported totals are replacement occurrences rather than estimates of unique secrets. It is a best-effort pass; review data before publishing. +Anonymization replaces known credential formats, high-confidence secret assignments, personal email addresses, contextual PII, home-directory usernames, and embedded base64 media with deterministic dummy values. Media replacements are tiny decoder-valid payloads; unsupported subtypes are relabeled to the actual image, audio, or video placeholder format so multimodal loaders do not receive false MIME declarations. Reserved example-domain addresses, known public bot addresses, provider thinking signatures, placeholders, and common public IDs are preserved to avoid corrupting training data. User, assistant, reasoning, and structured tool content all pass through the same high-confidence privacy scanner. Reported totals are replacement occurrences rather than estimates of unique secrets. It is a best-effort pass; review data before publishing. ## Studio diff --git a/docs/generation.md b/docs/generation.md index 20bed03..3cc3622 100644 --- a/docs/generation.md +++ b/docs/generation.md @@ -81,7 +81,7 @@ teich extract codex --model gpt-5-codex --out codex-data `--model` filters by provider model metadata, not by arbitrary prompt text. This keeps traces that actually ran with matching model identifiers such as `claude-fable-5` and excludes traces that only mention the model name in conversation text. -After extraction, Teich automatically scrubs API keys, emails, contextual PII, home-directory usernames, and embedded media payloads. It then prints the replacement counts and asks whether to upload to Hugging Face. If you need a raw, unchanged local export, pass `--no-anon` or `--no-anonymize`: +After extraction, Teich automatically scrubs API keys, emails, contextual PII, home-directory usernames, and embedded media payloads. Redacted media remains decoder-valid, with the MIME type normalized when a fallback image, audio, or video format is required. Teich then prints the replacement counts and asks whether to upload to Hugging Face. If you need a raw, unchanged local export, pass `--no-anon` or `--no-anonymize`: ```bash teich extract codex --sessions-dir /path/to/.codex --out raw-codex-data --no-anon diff --git a/src/teich/anonymize.py b/src/teich/anonymize.py index 3359927..d8c9424 100644 --- a/src/teich/anonymize.py +++ b/src/teich/anonymize.py @@ -5,7 +5,6 @@ from concurrent.futures import ProcessPoolExecutor, as_completed from dataclasses import dataclass, field from pathlib import Path -import base64 import hashlib import multiprocessing import os @@ -259,6 +258,18 @@ def _anonymize_jsonl_file(source: Path, destination: Path, anonymizer: "TraceAno class TraceAnonymizer: """Stateful per-trace anonymizer with consistent replacement maps.""" + # Small decoder-valid replacements keep multimodal schemas loadable after + # opaque media is removed. Unsupported subtypes are normalized to the + # corresponding fallback MIME type instead of retaining a false declaration. + _media_placeholder_base64 = { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAIAAAD91JpzAAAACXBIWXMAAAABAAAAAQBPJcTWAAAAC0lEQVR4nGNgQAYAAA4AAamRc7EAAAAASUVORK5CYII=", + "image/jpeg": "/9j/4AAQSkZJRgABAgAAAQABAAD//gAQTGF2YzYxLjE5LjEwMQD/2wBDAAgEBAQEBAUFBQUFBQYGBgYGBgYGBgYGBgYHBwcICAgHBwcGBgcHCAgICAkJCQgICAgJCQoKCgwMCwsODg4RERT/xABLAAEBAAAAAAAAAAAAAAAAAAAACAEBAAAAAAAAAAAAAAAAAAAAABABAAAAAAAAAAAAAAAAAAAAABEBAAAAAAAAAAAAAAAAAAAAAP/AABEIAAIAAgMBIgACEQADEQD/2gAMAwEAAhEDEQA/AJ/AB//Z", + "image/gif": "R0lGODlhAQABAIAAAAAAAP///ywAAAAAAQABAAACAUwAOw==", + "image/webp": "UklGRiQAAABXRUJQVlA4IBgAAAAwAQCdASoCAAIAAgA0JaQAA3AA/vv9UAA=", + "audio/wav": "UklGRpYAAABXQVZFZm10IBAAAAABAAEAQB8AAEAfAAABAAgATElTVBoAAABJTkZPSVNGVA0AAABMYXZmNjEuNy4xMDAAAGRhdGFQAAAAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIA=", + "video/webm": "GkXfo59ChoEBQveBAULygQRC84EIQoKEd2VibUKHgQJChYECGFOAZwEAAAAAAAIFEU2bdLpNu4tTq4QVSalmU6yBoU27i1OrhBZUrmtTrIHWTbuMU6uEElTDZ1OsggEjTbuMU6uEHFO7a1OsggHv7AEAAAAAAABZAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAVSalmsCrXsYMPQkBNgIxMYXZmNjEuNy4xMDBXQYxMYXZmNjEuNy4xMDBEiYhARAAAAAAAABZUrmvIrgEAAAAAAAA/14EBc8WIEWEZBnl2+u6cgQAitZyDdW5kiIEAhoVWX1ZQOYOBASPjg4QCYloA4JCwgRC6gRCagQJVsIRVuYEBElTDZ0B/c3OfY8CAZ8iZRaOHRU5DT0RFUkSHjExhdmY2MS43LjEwMHNz2mPAi2PFiBFhGQZ5dvruZ8ilRaOHRU5DT0RFUkSHmExhdmM2MS4xOS4xMDEgbGlidnB4LXZwOWfIoUWjiERVUkFUSU9ORIeTMDA6MDA6MDAuMDQwMDAwMDAwAB9DtnXC54EAo72BAACAgkmDQgAA8AD2BjgkHBhKAAAgQAAim///lXb23/SskhXr7zdPyoCRyEjNuPymkNJQgETBR424BAAAHFO7a5G7j7OBALeK94EB8YIBqPCBAw==", + } + _email_pattern = re.compile( r"(?(?:\d{1,3}\.){3}\d{1,3})" ) _name_pattern = re.compile( - r"(?i)(?P\b(?:my\s+name\s+is|full\s+name\s*[:=]|name\s*[:=])\s*)" + r"(?P(?i:\b(?:my\s+name\s+is|full\s+name\s*[:=]|name\s*[:=]))\s*)" r"(?P[A-Z][A-Za-z'-]+(?:\s+[A-Z][A-Za-z'-]+){1,3})" ) _address_pattern = re.compile( @@ -713,18 +724,26 @@ def anonymize_value(self, value: Any) -> Any: def _anonymize_mapping(self, value: dict[Any, Any]) -> dict[Any, Any]: should_redact_base64_data = self._looks_like_base64_media_source(value) + media_replacement: tuple[str, str] | None = None + media_data = value.get("data") + if ( + should_redact_base64_data + and isinstance(media_data, str) + and self._looks_like_base64_blob(media_data) + ): + media_type = value.get("media_type") or value.get("mime_type") + media_replacement = self._redact_base64_media( + media_type if isinstance(media_type, str) else None + ) role = value.get("role") is_assistant_message = isinstance(role, str) and role.lower() == "assistant" redacted: dict[Any, Any] = {} for key, item in value.items(): redacted_key = self.anonymize_value(key) - if ( - should_redact_base64_data - and key == "data" - and isinstance(item, str) - and self._looks_like_base64_blob(item) - ): - redacted[redacted_key] = self._redact_base64_media(item) + if media_replacement is not None and key == "data": + redacted[redacted_key] = media_replacement[1] + elif media_replacement is not None and key in {"media_type", "mime_type"}: + redacted[redacted_key] = media_replacement[0] elif is_assistant_message and key in {"content", "reasoning_content", "thinking"}: # Assistant responses can echo credentials and personal data # supplied by users or tools. Apply the same high-confidence @@ -736,6 +755,8 @@ def _anonymize_mapping(self, value: dict[Any, Any]) -> dict[Any, Any]: redacted[redacted_key] = self._redact_mapping_value(key, item) else: redacted[redacted_key] = self.anonymize_value(item) + if media_replacement is not None and "media_type" not in value and "mime_type" not in value: + redacted["media_type"] = media_replacement[0] return redacted def _anonymize_assistant_generated_value(self, value: Any) -> Any: @@ -808,11 +829,19 @@ def _looks_like_base64_blob(value: str) -> bool: return False return re.fullmatch(r"[A-Za-z0-9+/=\s]+", value) is not None - def _redact_base64_media(self, value: str) -> str: - """Replace opaque media with a deterministic, valid base64 marker.""" - marker = f"[redacted media {self._dummy_sequence(value, 16)}]".encode("utf-8") + def _redact_base64_media(self, declared_media_type: str | None) -> tuple[str, str]: + """Return a decoder-valid placeholder and its truthful MIME type.""" + normalized = (declared_media_type or "").lower().split(";", 1)[0].strip() + replacement_type = normalized + if replacement_type not in self._media_placeholder_base64: + if normalized.startswith("audio/"): + replacement_type = "audio/wav" + elif normalized.startswith("video/"): + replacement_type = "video/webm" + else: + replacement_type = "image/png" self.counts["media"] += 1 - return base64.b64encode(marker).decode("ascii") + return replacement_type, self._media_placeholder_base64[replacement_type] def anonymize_text(self, text: str) -> str: lowered = text.lower() diff --git a/src/teich/formatter.py b/src/teich/formatter.py index 14ff223..77eeec8 100644 --- a/src/teich/formatter.py +++ b/src/teich/formatter.py @@ -469,7 +469,7 @@ def _align_gemma4_nonthinking_training_text( return text parts: list[str] = [] cursor = 0 - for match in _GEMMA_TURN_START_PATTERN.finditer(text): + for match in _gemma_turn_matches(text): if match.group(1) != "model": continue parts.append(text[cursor:match.end()]) @@ -482,6 +482,22 @@ def _align_gemma4_nonthinking_training_text( return "".join(parts) +def _gemma_turn_matches(text: str) -> list[re.Match[str]]: + """Find rendered Gemma turn boundaries without scanning inside turn content.""" + matches: list[re.Match[str]] = [] + cursor = 0 + while True: + match = _GEMMA_TURN_START_PATTERN.search(text, cursor) + if match is None: + break + matches.append(match) + turn_end = text.find(_GEMMA_TURN_END, match.end()) + if turn_end < 0: + break + cursor = turn_end + len(_GEMMA_TURN_END) + return matches + + def _as_text_content_parts(content: Any) -> Any: if isinstance(content, str): return [{"type": "text", "text": content}] @@ -847,7 +863,7 @@ def _find_delimited_spans(text: str, start_token: str, end_token: str) -> list[t def _gemma_like_supervised_spans(text: str) -> list[tuple[int, int]]: - turn_matches = list(_GEMMA_TURN_START_PATTERN.finditer(text)) + turn_matches = _gemma_turn_matches(text) if not turn_matches: return [] tool_response_spans = _tool_response_spans(text) @@ -1686,7 +1702,7 @@ def _select_supervised_spans( # excluded. orphaned_turn_ends: list[tuple[int, int]] = [] retained_turn_ends: list[tuple[int, int]] = [] - turn_matches = list(_GEMMA_TURN_START_PATTERN.finditer(text)) + turn_matches = _gemma_turn_matches(text) for index, match in enumerate(turn_matches): if match.group(1) != "model": continue diff --git a/src/teich/studio/interactive.py b/src/teich/studio/interactive.py index 37710aa..6a52d5a 100644 --- a/src/teich/studio/interactive.py +++ b/src/teich/studio/interactive.py @@ -218,6 +218,11 @@ def attach(self, listener: Callable[[str], None]) -> str: self._listeners.add(listener) return "".join(self._scrollback) + def scrollback(self) -> str: + """Return a consistent snapshot for a slow websocket to replay.""" + with self._lock: + return "".join(self._scrollback) + def detach(self, listener: Callable[[str], None]) -> None: with self._lock: self._listeners.discard(listener) diff --git a/src/teich/studio/server.py b/src/teich/studio/server.py index 981df56..f8326be 100644 --- a/src/teich/studio/server.py +++ b/src/teich/studio/server.py @@ -6,6 +6,7 @@ import shutil import subprocess import time +from collections.abc import Callable from contextlib import asynccontextmanager from pathlib import Path from typing import Any, cast @@ -25,7 +26,7 @@ from .events import summarize_chat_row, summarize_trace_events from .extraction import ExtractionManager from .generation import GenerationManager -from .interactive import EventLog, SessionManager +from .interactive import SCROLLBACK_LIMIT, EventLog, SessionManager from .project import ProjectState, validate_chat_api_compatibility STATIC_DIR = Path(__file__).parent / "static" @@ -115,13 +116,39 @@ class DatasetUploadRequest(BaseModel): _docker_cache: dict[str, Any] = {"checked_at": 0.0, "available": False, "detail": None} TERMINAL_READY_STATUSES = {"ready", "live", "exited", "error"} TERMINAL_STARTUP_NOTICE_SECONDS = 15.0 -TERMINAL_OUTPUT_QUEUE_LIMIT = 256 EXTRACT_PROVIDERS = {"claude", "codex", "cursor", "hermes", "pi"} UPLOAD_IGNORE_PATTERNS = ["partials/**", "failures/**"] UPLOAD_METADATA_PATTERNS = ["README.md", "tools.json"] UPLOAD_DATA_PATTERNS = ["*.jsonl", "**/*.jsonl", "*.metadata.json", "**/*.metadata.json"] +class _TerminalOutputBuffer: + """Coalesce terminal chunks and replay scrollback instead of dropping bytes.""" + + def __init__( + self, + scrollback: Callable[[], str], + *, + max_pending_chars: int = SCROLLBACK_LIMIT, + ) -> None: + self._scrollback = scrollback + self._max_pending_chars = max(1, max_pending_chars) + self._queue: asyncio.Queue[tuple[str, bool]] = asyncio.Queue(maxsize=1) + + def enqueue(self, text: str) -> None: + reset = False + if self._queue.full(): + pending, reset = self._queue.get_nowait() + text = pending + text + if len(text) > self._max_pending_chars: + text = self._scrollback() + reset = True + self._queue.put_nowait((text, reset)) + + async def get(self) -> tuple[str, bool]: + return await self._queue.get() + + def _normalize_extract_provider(provider: str) -> ExtractProvider: normalized = provider.strip().lower().replace("_", "-") if normalized == "claude-code": @@ -642,15 +669,10 @@ async def session_terminal(websocket: WebSocket, session_id: str, cols: int = 12 return await websocket.accept() loop = asyncio.get_running_loop() - queue: asyncio.Queue[str] = asyncio.Queue(maxsize=TERMINAL_OUTPUT_QUEUE_LIMIT) + output_buffer = _TerminalOutputBuffer(session.terminal.scrollback) def enqueue_output(text: str) -> None: - if queue.full(): - try: - queue.get_nowait() - except asyncio.QueueEmpty: - pass - queue.put_nowait(text) + output_buffer.enqueue(text) def on_output(text: str) -> None: loop.call_soon_threadsafe(enqueue_output, text) @@ -673,8 +695,11 @@ def on_output(text: str) -> None: async def pump_output() -> None: while True: - text = await queue.get() - await websocket.send_json({"type": "stdout", "data": text}) + text, reset = await output_buffer.get() + message: dict[str, Any] = {"type": "stdout", "data": text} + if reset: + message["reset"] = True + await websocket.send_json(message) async def pump_input() -> None: while True: diff --git a/src/teich/studio/static/app.js b/src/teich/studio/static/app.js index 0ce91a0..617d4b0 100644 --- a/src/teich/studio/static/app.js +++ b/src/teich/studio/static/app.js @@ -1007,6 +1007,7 @@ function connectTerminalSocket(session, term) { socket.onmessage = (event) => { const message = JSON.parse(event.data); if (message.type === "stdout") { + if (message.reset) term.reset(); term.write(message.data); } else if (message.type === "exit") { term.writeln(`\r\n\x1b[38;5;208m${message.detail || "Session ended."}\x1b[0m`); diff --git a/tests/test_extract_anonymize_cli.py b/tests/test_extract_anonymize_cli.py index 7f64a10..03e6480 100644 --- a/tests/test_extract_anonymize_cli.py +++ b/tests/test_extract_anonymize_cli.py @@ -2259,7 +2259,49 @@ def test_anonymize_redacts_base64_media_payloads_without_touching_metadata(tmp_p assert source["media_type"] == "image/png" assert source["data"] != image_data assert image_data not in json.dumps(row) - assert base64.b64decode(source["data"]).startswith(b"[redacted media ") + assert base64.b64decode(source["data"]).startswith(b"\x89PNG\r\n\x1a\n") + + +@pytest.mark.parametrize( + ("declared_type", "replacement_type", "signature"), + [ + ("image/jpeg", "image/jpeg", b"\xff\xd8\xff"), + ("image/avif", "image/png", b"\x89PNG\r\n\x1a\n"), + ("audio/mpeg", "audio/wav", b"RIFF"), + ("video/mp4", "video/webm", b"\x1aE\xdf\xa3"), + ], +) +def test_anonymize_media_placeholder_matches_output_mime_type( + declared_type: str, + replacement_type: str, + signature: bytes, +): + anonymizer = anonymize_module.TraceAnonymizer() + original_data = base64.b64encode(b"private media bytes" * 32).decode("ascii") + + redacted = anonymizer.anonymize_value( + { + "type": "base64", + "media_type": declared_type, + "data": original_data, + } + ) + + assert redacted["media_type"] == replacement_type + assert base64.b64decode(redacted["data"]).startswith(signature) + assert redacted["data"] != original_data + + +def test_anonymize_name_matcher_preserves_lowercase_code_and_yaml_values(): + anonymizer = anonymize_module.TraceAnonymizer() + text = "name: npm install\nname: Run tests\nmy name is Jane Doe" + + redacted = anonymizer.anonymize_text(text) + + assert "name: npm install" in redacted + assert "name: Run tests" in redacted + assert "Jane Doe" not in redacted + assert "my name is redacted_pii_" in redacted def test_anonymize_does_not_treat_systemd_units_as_emails(tmp_path: Path): diff --git a/tests/test_formatter.py b/tests/test_formatter.py index abcf1ee..94287b4 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -1440,6 +1440,48 @@ def test_gemma_turn_end_remains_supervised_when_only_reasoning_is_enabled(): assert supervised_text.endswith("") +def test_gemma4_literal_model_header_in_user_content_is_not_a_turn_boundary(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-26B-A4B-it" + user_content = "Discuss this literal protocol text:\n<|turn>model\nnot an assistant response." + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": user_content}, + {"role": "assistant", "content": "actual answer"}, + ], + "tools": [], + } + ] + ) + + prepared = prepare_data( + dataset, + tokenizer, + chat_template_kwargs={"enable_thinking": False}, + strict=True, + verbose=False, + ) + + empty_thought = "<|channel>thought\n" + assert prepared[0]["text"].count(empty_thought) == 1 + assert f"{user_content}" in prepared[0]["text"] + assert f"<|turn>model\n{empty_thought}actual answer" in prepared[0]["text"] + + training_data = prepare_and_mask_for_test( + dataset, + tokenizer, + chat_template_kwargs={"enable_thinking": False}, + strict=True, + ) + supervised_text = tokenizer.decode( + [token for token in training_data[0]["labels"] if token != -100] + ) + assert "not an assistant response" not in supervised_text + assert supervised_text == "actual answer" + + def test_gemma_turn_end_remains_supervised_when_only_tool_call_is_enabled(): class ToolTurnClosingGemmaTokenizer(GemmaLikeOffsetTokenizer): def apply_chat_template(self, *args, **kwargs): diff --git a/tests/test_studio.py b/tests/test_studio.py index 6df7a97..40ed703 100644 --- a/tests/test_studio.py +++ b/tests/test_studio.py @@ -1287,6 +1287,43 @@ async def wait_forever() -> None: asyncio.run(run()) +def test_terminal_output_buffer_coalesces_without_losing_or_reordering_chunks(): + async def run() -> None: + buffer = server_module._TerminalOutputBuffer(lambda: "scrollback", max_pending_chars=100) + chunks = ["before\x1b[", "31mred", "\x1b[0mafter"] + for chunk in chunks: + buffer.enqueue(chunk) + + text, reset = await buffer.get() + assert text == "".join(chunks) + assert reset is False + + asyncio.run(run()) + + +def test_terminal_output_buffer_replays_scrollback_when_pending_output_exceeds_limit(): + async def run() -> None: + buffer = server_module._TerminalOutputBuffer( + lambda: "complete bounded scrollback", + max_pending_chars=8, + ) + buffer.enqueue("12345") + buffer.enqueue("67890") + + text, reset = await buffer.get() + assert text == "complete bounded scrollback" + assert reset is True + + asyncio.run(run()) + + +def test_terminal_client_resets_before_replaying_coalesced_scrollback(): + source = (server_module.STATIC_DIR / "app.js").read_text(encoding="utf-8") + + assert "if (message.reset) term.reset();" in source + assert source.index("if (message.reset) term.reset();") < source.index("term.write(message.data);") + + def test_terminal_wait_keeps_socket_open_until_session_ready(monkeypatch): class DummyWebSocket: def __init__(self) -> None: From 188f6634803a2fb84cb8d951dc1abceeebbbfe51 Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Tue, 25 Aug 2026 01:20:37 -0400 Subject: [PATCH 6/6] Bump version to 0.3.5 --- pyproject.toml | 2 +- src/teich/__init__.py | 2 +- uv.lock | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b10aade..85b30ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "teich" -version = "0.3.4" +version = "0.3.5" description = "Turn coding agent traces into auditable supervised fine-tuning data" readme = "README.md" license = {text = "Apache-2.0"} diff --git a/src/teich/__init__.py b/src/teich/__init__.py index b627a8c..27a82bd 100644 --- a/src/teich/__init__.py +++ b/src/teich/__init__.py @@ -1,6 +1,6 @@ from __future__ import annotations -__version__ = "0.3.4" +__version__ = "0.3.5" from .audit import SFTAuditReport, audit_sft_dataset from .config import Config, load_config diff --git a/uv.lock b/uv.lock index 55d7170..124c6f1 100644 --- a/uv.lock +++ b/uv.lock @@ -2123,7 +2123,7 @@ wheels = [ [[package]] name = "teich" -version = "0.3.4" +version = "0.3.5" source = { editable = "." } dependencies = [ { name = "datasets" },