trainsight-preflight 0.2.0__py3-none-any.whl

trainsight/__init__.py

@@ -0,0 +1,39 @@
+"""
+TrainSight: Preflight Checks for LLM Training Inputs and Loss Contracts.
+"""
+
+__version__ = "0.2.0"
+
+from trainsight.core.contracts import LossContract
+from trainsight.core.exceptions import (
+    TrainSightError,
+    ContractViolationError,
+    PreflightQualityGateError,
+)
+from trainsight.core.report import AuditReport, RowDiagnosis, TensorEvidence
+from trainsight.inspectors.sft_inspector import SFTInspector, SFTReport
+from trainsight.inspectors.dpo_inspector import DPOInspector, DPOReport
+from trainsight.preprocessing.sft_preprocessor import SFTPreprocessor, SFTExample
+from trainsight.preprocessing.dpo_preprocessor import DPOPreprocessor, DPOExample
+from trainsight.reconcilers.sft_reconciler import SFTReconciler
+from trainsight.reconcilers.callbacks import TrainSightCallback
+
+__all__ = [
+    "LossContract",
+    "TrainSightError",
+    "ContractViolationError",
+    "PreflightQualityGateError",
+    "AuditReport",
+    "RowDiagnosis",
+    "TensorEvidence",
+    "SFTInspector",
+    "SFTReport",
+    "DPOInspector",
+    "DPOReport",
+    "SFTPreprocessor",
+    "SFTExample",
+    "DPOPreprocessor",
+    "DPOExample",
+    "SFTReconciler",
+    "TrainSightCallback",
+]

trainsight/cli/__init__.py

@@ -0,0 +1,3 @@
+from trainsight.cli.main import app
+
+__all__ = ["app"]

trainsight/cli/check.py

@@ -0,0 +1,80 @@
+from pathlib import Path
+from typing import Optional
+import typer
+from rich.table import Table
+
+from trainsight.inspectors.sft_inspector import SFTInspector
+from trainsight.inspectors.dpo_inspector import DPOInspector
+from trainsight.utils.logging import get_console
+
+app = typer.Typer(help="Inspect raw training datasets on disk before GPU allocation.")
+
+@app.callback(invoke_without_command=True)
+def check(
+    file_path: Path = typer.Argument(..., help="Path to JSONL dataset file"),
+    task: str = typer.Option("sft", "--task", "-t", help="Task type: sft or dpo"),
+    tokenizer: str = typer.Option("Qwen/Qwen2.5-0.5B-Instruct", "--tokenizer", "-m", help="Tokenizer name or path"),
+    max_seq_len: int = typer.Option(2048, "--max-seq-len", "-l", help="Max sequence length"),
+    strict: bool = typer.Option(True, "--strict/--no-strict", help="Fail with non-zero exit code on contract violations"),
+):
+    console = get_console()
+    if not file_path.exists():
+        console.print(f"[bold red]Error:[/] File not found: {file_path}")
+        raise typer.Exit(code=1)
+
+    console.print(f"\n[bold cyan]TrainSight Preflight Inspection[/]")
+    console.print(f"File: [bold]{file_path.name}[/] | Task: [bold]{task.upper()}[/] | Max Length: {max_seq_len}\n")
+
+    if task.lower() == "sft":
+        inspector = SFTInspector(
+            tokenizer_name=tokenizer,
+            max_seq_len_threshold=max_seq_len,
+            strict=strict,
+        )
+        report = inspector.inspect_file(file_path)
+
+        table = Table(title="SFT Preflight Summary", show_lines=True)
+        table.add_column("Metric", style="cyan")
+        table.add_column("Value", justify="right")
+        table.add_column("Status", justify="left")
+
+        table.add_row("Total Samples", str(report.total_samples), "[green]OK[/]")
+        table.add_row("Valid Samples", str(report.valid_samples_count), "[green]OK[/]")
+        
+        status_empty = "[green]OK[/]" if report.empty_completion_count == 0 else "[bold red]FAIL[/]"
+        table.add_row("Empty Completions", str(report.empty_completion_count), status_empty)
+
+        status_trunc = "[green]OK[/]" if report.truncated_answers_count == 0 else "[bold red]FAIL[/]"
+        table.add_row("Truncated Answers", str(report.truncated_answers_count), status_trunc)
+
+        status_mask = "[green]OK[/]" if report.loss_mask_misaligned_count == 0 else "[bold red]FAIL[/]"
+        table.add_row("Misaligned Loss Masks", str(report.loss_mask_misaligned_count), status_mask)
+
+        table.add_row("Mean Sequence Length", f"{report.avg_seq_len:.1f} tokens", "[dim]Info[/]")
+        table.add_row("Max Sequence Length", f"{report.max_seq_len} tokens", "[dim]Info[/]")
+
+        console.print(table)
+
+        passed = (
+            report.invalid_rows_count == 0
+            and report.empty_completion_count == 0
+            and report.truncated_answers_count == 0
+            and report.loss_mask_misaligned_count == 0
+        )
+
+        if passed:
+            console.print("\n[bold green]✓ Preflight Quality Gate Passed:[/] Dataset meets all loss contract requirements.\n")
+        else:
+            console.print("\n[bold red]✗ Preflight Quality Gate Failed:[/] Contract violations detected above.\n")
+            if strict:
+                raise typer.Exit(code=1)
+
+    elif task.lower() == "dpo":
+        inspector = DPOInspector(
+            tokenizer_name=tokenizer,
+            max_seq_len_threshold=max_seq_len,
+        )
+        report = inspector.inspect_file(file_path)
+        console.print(f"Total pairs: {report.total_pairs} | Valid pairs: {report.valid_pairs_count}")
+        if report.invalid_pairs_count > 0 and strict:
+            raise typer.Exit(code=1)

trainsight/cli/demo.py

@@ -0,0 +1,131 @@
+import json
+import tempfile
+from pathlib import Path
+import typer
+from rich.panel import Panel
+import torch
+from transformers import AutoTokenizer
+
+from trainsight.inspectors.sft_inspector import SFTInspector
+from trainsight.core.contracts import LossContract
+from trainsight.core.exceptions import ContractViolationError
+from trainsight.reconcilers.sft_reconciler import SFTReconciler
+from trainsight.utils.logging import get_console
+
+app = typer.Typer(help="Run interactive demonstration of preflight and batch checks.")
+
+DEMO_CLEAN = [
+    {"id": "row_001", "prompt": "What is the capital of France?", "completion": "The capital of France is Paris."},
+    {"id": "row_002", "prompt": "Solve: 15 * 8", "completion": "15 * 8 = 120."},
+    {"id": "row_003", "prompt": "Is 37 a prime number?", "completion": "Yes, 37 is a prime number."},
+]
+
+@app.callback(invoke_without_command=True)
+def run_demo():
+    console = get_console()
+    console.print("\n[bold cyan]========================================================================[/]")
+    console.print("[bold cyan]       TrainSight: Preflight LLM Training Input & Loss Contract Auditor  [/]")
+    console.print("[bold cyan]========================================================================[/]\n")
+    console.print("[dim]Note: First run downloads tokenizer configuration files if not already cached.[/]\n")
+
+    model_name = "Qwen/Qwen2.5-0.5B-Instruct"
+    console.print(f"[dim]Loading tokenizer ({model_name})...[/]")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    except Exception as e:
+        console.print(f"[bold red]Failed to load tokenizer:[/] {e}")
+        raise typer.Exit(code=1)
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # -------------------------------------------------------------
+        # STEP 1: Clean Input Verification (Validated against actual report)
+        # -------------------------------------------------------------
+        console.print("[bold]Step 1: Inspecting clean synthetic dataset (3 rows)...[/]")
+        clean_file = tmp_path / "clean.jsonl"
+        with open(clean_file, "w", encoding="utf-8") as f:
+            for r in DEMO_CLEAN:
+                f.write(json.dumps(r) + "\n")
+
+        inspector = SFTInspector(tokenizer_name=model_name, max_seq_len_threshold=256)
+        clean_report = inspector.inspect_file(clean_file)
+
+        console.print(f"  • Total samples: {clean_report.total_samples} | Valid: {clean_report.valid_samples_count}")
+        console.print(f"  • Empty completions: {clean_report.empty_completion_count} | Truncated answers: {clean_report.truncated_answers_count}")
+
+        # Check actual report before declaring pass
+        clean_passed = (
+            clean_report.valid_samples_count == clean_report.total_samples
+            and clean_report.empty_completion_count == 0
+            and clean_report.truncated_answers_count == 0
+            and clean_report.loss_mask_misaligned_count == 0
+        )
+        if clean_passed:
+            console.print("[bold green]✓ Preflight Quality Gate Passed (Zero False Alarms)[/]\n")
+        else:
+            console.print("[bold red]✗ Unexpected Preflight Failure on Clean Data[/]\n")
+
+        # -------------------------------------------------------------
+        # STEP 2: Truncation Elimination Fault (Validated against actual report)
+        # -------------------------------------------------------------
+        console.print("[bold]Step 2: Injecting deliberate truncation fault (Answer eliminated by budget)...[/]")
+        fault_file = tmp_path / "fault_truncation.jsonl"
+        long_prompt = "Verbose prompt designed to consume token budget. " * 30
+        with open(fault_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps({"id": "row_004", "prompt": long_prompt, "completion": "Answer that will be clipped."}) + "\n")
+
+        fault_report = inspector.inspect_file(fault_file)
+        if fault_report.truncated_answers_count > 0:
+            console.print("[bold red]✗ Preflight Quality Gate Caught Fault from Report Results:[/]")
+            console.print(Panel(
+                f"[bold red]CRITICAL CONTRACT VIOLATION:[/] {fault_report.truncated_answers_count} answer(s) eliminated by truncation.\\n"
+                f"Row ID: [bold]row_004[/]\\n"
+                f"• Prompt consumed available token budget (max_seq_len=256).\\n"
+                f"• Retained answer tokens: 0.\\n"
+                f"• Impact: Example retains zero supervised targets.",
+                title="TrainSight Preflight Gate (Actual Inspection Report)",
+                border_style="red",
+            ))
+        else:
+            console.print("[bold yellow]! Preflight Gate did not flag truncation on fault file.[/]")
+
+        # -------------------------------------------------------------
+        # STEP 3: Synthetic Corrupted Tensor Batch Audit
+        # -------------------------------------------------------------
+        console.print("\n[bold]Step 3: Auditing a synthetic corrupted tensor batch (prompt unmasking leak)...[/]")
+        p_text = tokenizer.apply_chat_template([{"role": "user", "content": "Hello"}], tokenize=False, add_generation_prompt=True)
+        c_text = "Hi there!<|im_end|>"
+        p_ids = tokenizer.encode(p_text, add_special_tokens=False)
+        c_ids = tokenizer.encode(c_text, add_special_tokens=False)
+        full_ids = p_ids + c_ids
+
+        # Construct a synthetic corrupted collator output: prompt unmasked (labels[:p_len] != -100)
+        corrupted_batch = {
+            "input_ids": torch.tensor([full_ids]),
+            "labels": torch.tensor([full_ids]),  # Corrupted: prompt labels equal token IDs instead of -100
+            "attention_mask": torch.tensor([[1] * len(full_ids)]),
+        }
+
+        reconciler = SFTReconciler(tokenizer=tokenizer, contract=LossContract(strict=False))
+        audit = reconciler.audit_batch(batch=corrupted_batch, prompt_lengths=[len(p_ids)], row_ids=["row_synthetic_0"])
+
+        if not audit.passed and len(audit.diagnoses) > 0:
+            diag = audit.diagnoses[0]
+            console.print("[bold red]✗ Batch Reconciler Caught Post-Collation Mask Leakage:[/]")
+            console.print(Panel(
+                f"[bold red]COLLATOR LEAKAGE:[/] {diag.error_code}\\n"
+                f"{diag.error_message}\\n"
+                f"• Prompt length: {diag.evidence.prompt_token_length} tokens.\\n"
+                f"• Unmasked prompt tokens receiving loss: {diag.evidence.unmasked_prompt_tokens_count}.\\n"
+                f"• Impact: Gradients would be computed over user prompt tokens.",
+                title="TrainSight Batch Reconciler (Synthetic Corrupted Batch)",
+                border_style="red",
+            ))
+        else:
+            console.print("[bold green]✓ Reconciler passed synthetic batch.[/]")
+
+    console.print("\n[bold green]========================================================================[/]")
+    console.print("[bold green]✓ Demo Complete: Verified preflight gate and batch reconciler diagnostics.[/]")
+    console.print("[bold green]========================================================================[/]\n")
+    console.print("To inspect your own dataset:\n  $ [bold cyan]trainsight check my_data.jsonl --tokenizer Qwen/Qwen2.5-0.5B-Instruct[/]\n")

trainsight/cli/main.py

@@ -0,0 +1,22 @@
+import typer
+from trainsight import __version__
+from trainsight.cli.check import check
+from trainsight.cli.demo import run_demo
+from trainsight.utils.logging import get_console
+
+app = typer.Typer(
+    name="trainsight",
+    help="Preflight checks for LLM training inputs and loss contracts.",
+    add_completion=False,
+)
+
+app.command(name="check")(check)
+app.command(name="demo")(run_demo)
+
+@app.command(name="version")
+def version():
+    """Prints TrainSight version."""
+    get_console().print(f"TrainSight version: [bold cyan]{__version__}[/]")
+
+if __name__ == "__main__":
+    app()

trainsight/core/__init__.py

@@ -0,0 +1,13 @@
+from trainsight.core.contracts import LossContract
+from trainsight.core.exceptions import TrainSightError, ContractViolationError, PreflightQualityGateError
+from trainsight.core.report import AuditReport, RowDiagnosis, TensorEvidence
+
+__all__ = [
+    "LossContract",
+    "TrainSightError",
+    "ContractViolationError",
+    "PreflightQualityGateError",
+    "AuditReport",
+    "RowDiagnosis",
+    "TensorEvidence",
+]

trainsight/core/contracts.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Optional, Literal
+
+@dataclass
+class LossContract:
+    """Declared contract for loss calculation on training batches."""
+    task_type: Literal["sft", "dpo"] = "sft"
+    chat_template: str = "chatml"
+    max_seq_len: int = 2048
+    answer_only_loss: bool = True
+    mask_padding: bool = True
+    require_causal_shift_supervision: bool = True
+    strict: bool = True
+    expected_template_prefix: str = "<|im_start|>user"
+    expected_template_suffix: str = "<|im_start|>assistant\n"

trainsight/core/exceptions.py

@@ -0,0 +1,14 @@
+class TrainSightError(Exception):
+    """Base exception for all TrainSight errors."""
+    pass
+
+class ContractViolationError(TrainSightError):
+    """Raised when actual training tensors violate the intended loss contract."""
+    def __init__(self, check_name: str, message: str, row_id: str = "unknown"):
+        self.check_name = check_name
+        self.row_id = row_id
+        super().__init__(f"[{check_name}] (Row: {row_id}): {message}")
+
+class PreflightQualityGateError(TrainSightError):
+    """Raised when raw dataset inspection fails a pre-flight quality gate."""
+    pass

trainsight/core/report.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass, field
+from typing import List, Dict, Any, Optional
+
+@dataclass
+class TensorEvidence:
+    row_id: str
+    sequence_length: int
+    prompt_token_length: int
+    active_labels_count: int
+    causal_loss_contributing_tokens: int
+    unmasked_prompt_tokens_count: int
+    unmasked_padding_tokens_count: int
+    decoded_prompt_preview: str
+    decoded_answer_preview: str
+    labels_sample: List[Any]
+
+@dataclass
+class RowDiagnosis:
+    row_id: str
+    error_code: str
+    error_message: str
+    severity: str = "critical"  # "critical" or "warning"
+    evidence: Optional[TensorEvidence] = None
+
+@dataclass
+class AuditReport:
+    total_inspected: int
+    valid_count: int
+    violations_count: int
+    passed: bool
+    diagnoses: List[RowDiagnosis] = field(default_factory=list)
+
+    def summary_dict(self) -> Dict[str, Any]:
+        return {
+            "total_inspected": self.total_inspected,
+            "valid_count": self.valid_count,
+            "violations_count": self.violations_count,
+            "passed": self.passed,
+            "diagnoses": [
+                {
+                    "row_id": d.row_id,
+                    "error_code": d.error_code,
+                    "error_message": d.error_message,
+                    "severity": d.severity,
+                }
+                for d in self.diagnoses
+            ],
+        }

trainsight/inspectors/__init__.py

@@ -0,0 +1,5 @@
+from trainsight.inspectors.base import BaseInspector
+from trainsight.inspectors.sft_inspector import SFTInspector, SFTReport
+from trainsight.inspectors.dpo_inspector import DPOInspector, DPOReport
+
+__all__ = ["BaseInspector", "SFTInspector", "SFTReport", "DPOInspector", "DPOReport"]

trainsight/inspectors/base.py

@@ -0,0 +1,9 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+class BaseInspector(ABC):
+    """Abstract base class for dataset inspectors."""
+    @abstractmethod
+    def inspect_file(self, file_path: Path) -> Any:
+        pass