synkro 0.4.12__py3-none-any.whl

synkro/core/checkpoint.py

@@ -0,0 +1,250 @@
+"""Checkpoint manager for resumable generation."""
+
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+if TYPE_CHECKING:
+    from synkro.types.core import Trace
+    from synkro.types.logic_map import LogicMap, GoldenScenario
+
+console = Console()
+
+
+class CheckpointData(BaseModel):
+    """Data stored in a checkpoint file."""
+
+    # Metadata
+    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())
+    policy_hash: str = ""
+    target_traces: int = 0
+    dataset_type: str = ""
+
+    # Stage 1: Logic Map
+    logic_map_data: dict | None = None
+
+    # Stage 2: Scenarios
+    scenarios_data: list[dict] = Field(default_factory=list)
+    scenario_distribution: dict[str, int] = Field(default_factory=dict)
+
+    # Stage 3: Generated traces
+    traces_data: list[dict] = Field(default_factory=list)
+    completed_scenario_indices: list[int] = Field(default_factory=list)
+
+    # Stage 4: Verified traces
+    verified_traces_data: list[dict] = Field(default_factory=list)
+    verification_complete: bool = False
+
+
+class CheckpointManager:
+    """
+    Manages checkpoints for resumable generation.
+
+    Saves progress after each stage and allows resuming from the last
+    successful checkpoint.
+
+    Examples:
+        >>> manager = CheckpointManager("./checkpoints")
+        >>> manager.save_logic_map(logic_map, policy_hash, 100, "sft")
+        >>> manager.save_scenarios(scenarios, distribution)
+        >>> manager.save_trace(trace, scenario_index)
+
+        >>> # Resume from checkpoint
+        >>> checkpoint = manager.load()
+        >>> if checkpoint:
+        ...     logic_map = checkpoint.get_logic_map()
+        ...     completed = checkpoint.completed_scenario_indices
+    """
+
+    def __init__(self, checkpoint_dir: str | Path):
+        """
+        Initialize the checkpoint manager.
+
+        Args:
+            checkpoint_dir: Directory to store checkpoint files
+        """
+        self.checkpoint_dir = Path(checkpoint_dir)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        self.checkpoint_file = self.checkpoint_dir / "checkpoint.json"
+        self._data: CheckpointData | None = None
+
+    def _load_or_create(self) -> CheckpointData:
+        """Load existing checkpoint or create new one."""
+        if self._data is not None:
+            return self._data
+
+        if self.checkpoint_file.exists():
+            with open(self.checkpoint_file) as f:
+                data = json.load(f)
+            self._data = CheckpointData.model_validate(data)
+        else:
+            self._data = CheckpointData()
+
+        return self._data
+
+    def _save(self) -> None:
+        """Save checkpoint to disk."""
+        if self._data is None:
+            return
+
+        with open(self.checkpoint_file, "w") as f:
+            json.dump(self._data.model_dump(), f, indent=2)
+
+    def has_checkpoint(self) -> bool:
+        """Check if a checkpoint exists."""
+        return self.checkpoint_file.exists()
+
+    def load(self) -> CheckpointData | None:
+        """Load checkpoint if it exists."""
+        if not self.has_checkpoint():
+            return None
+        return self._load_or_create()
+
+    def matches_config(self, policy_hash: str, target_traces: int, dataset_type: str) -> bool:
+        """Check if checkpoint matches the current generation config."""
+        data = self._load_or_create()
+        return (
+            data.policy_hash == policy_hash
+            and data.target_traces == target_traces
+            and data.dataset_type == dataset_type
+        )
+
+    def save_logic_map(
+        self,
+        logic_map: "LogicMap",
+        policy_hash: str,
+        target_traces: int,
+        dataset_type: str,
+    ) -> None:
+        """Save Logic Map (Stage 1 complete)."""
+        data = self._load_or_create()
+        data.policy_hash = policy_hash
+        data.target_traces = target_traces
+        data.dataset_type = dataset_type
+        data.logic_map_data = logic_map.model_dump()
+        self._save()
+        console.print("[dim]💾 Checkpoint: Logic Map saved[/dim]")
+
+    def save_scenarios(
+        self,
+        scenarios: list["GoldenScenario"],
+        distribution: dict[str, int],
+    ) -> None:
+        """Save scenarios (Stage 2 complete)."""
+        data = self._load_or_create()
+        data.scenarios_data = [s.model_dump() for s in scenarios]
+        data.scenario_distribution = distribution
+        self._save()
+        console.print("[dim]💾 Checkpoint: Scenarios saved[/dim]")
+
+    def save_trace(self, trace: "Trace", scenario_index: int) -> None:
+        """Save a generated trace (incremental Stage 3)."""
+        data = self._load_or_create()
+        data.traces_data.append(trace.model_dump())
+        data.completed_scenario_indices.append(scenario_index)
+        self._save()
+
+    def save_traces_batch(self, traces: list["Trace"], indices: list[int]) -> None:
+        """Save a batch of traces at once."""
+        data = self._load_or_create()
+        for trace, idx in zip(traces, indices):
+            data.traces_data.append(trace.model_dump())
+            data.completed_scenario_indices.append(idx)
+        self._save()
+        console.print(f"[dim]💾 Checkpoint: {len(traces)} traces saved[/dim]")
+
+    def save_verified_traces(self, traces: list["Trace"]) -> None:
+        """Save verified traces (Stage 4 complete)."""
+        data = self._load_or_create()
+        data.verified_traces_data = [t.model_dump() for t in traces]
+        data.verification_complete = True
+        self._save()
+        console.print("[dim]💾 Checkpoint: Verification complete[/dim]")
+
+    def get_logic_map(self) -> "LogicMap | None":
+        """Retrieve Logic Map from checkpoint."""
+        from synkro.types.logic_map import LogicMap
+
+        data = self._load_or_create()
+        if data.logic_map_data:
+            return LogicMap.model_validate(data.logic_map_data)
+        return None
+
+    def get_scenarios(self) -> list["GoldenScenario"]:
+        """Retrieve scenarios from checkpoint."""
+        from synkro.types.logic_map import GoldenScenario
+
+        data = self._load_or_create()
+        return [GoldenScenario.model_validate(s) for s in data.scenarios_data]
+
+    def get_traces(self) -> list["Trace"]:
+        """Retrieve traces from checkpoint."""
+        from synkro.types.core import Trace
+
+        data = self._load_or_create()
+        return [Trace.model_validate(t) for t in data.traces_data]
+
+    def get_verified_traces(self) -> list["Trace"]:
+        """Retrieve verified traces from checkpoint."""
+        from synkro.types.core import Trace
+
+        data = self._load_or_create()
+        return [Trace.model_validate(t) for t in data.verified_traces_data]
+
+    def get_pending_scenario_indices(self, total: int) -> list[int]:
+        """Get indices of scenarios that haven't been processed yet."""
+        data = self._load_or_create()
+        completed = set(data.completed_scenario_indices)
+        return [i for i in range(total) if i not in completed]
+
+    def clear(self) -> None:
+        """Clear the checkpoint."""
+        if self.checkpoint_file.exists():
+            self.checkpoint_file.unlink()
+        self._data = None
+        console.print("[dim]🗑️ Checkpoint cleared[/dim]")
+
+    @property
+    def stage(self) -> str:
+        """Get the current stage based on checkpoint data."""
+        data = self._load_or_create()
+
+        if data.verification_complete:
+            return "complete"
+        if data.traces_data:
+            return "traces"  # In progress or done
+        if data.scenarios_data:
+            return "scenarios"
+        if data.logic_map_data:
+            return "logic_map"
+        return "start"
+
+    def summary(self) -> str:
+        """Get a summary of the checkpoint status."""
+        data = self._load_or_create()
+
+        lines = [
+            f"Checkpoint Status",
+            f"=================",
+            f"Stage: {self.stage}",
+            f"Target traces: {data.target_traces}",
+            f"Logic Map: {'✓' if data.logic_map_data else '✗'}",
+            f"Scenarios: {len(data.scenarios_data)}",
+            f"Traces: {len(data.traces_data)}/{data.target_traces}",
+            f"Verified: {'✓' if data.verification_complete else '✗'}",
+        ]
+
+        return "\n".join(lines)
+
+
+def hash_policy(policy_text: str) -> str:
+    """Create a hash of policy text for checkpoint matching."""
+    import hashlib
+    return hashlib.sha256(policy_text.encode()).hexdigest()[:16]
+
+
+__all__ = ["CheckpointManager", "CheckpointData", "hash_policy"]

synkro/core/dataset.py

@@ -0,0 +1,402 @@
+"""Dataset class for managing generated traces."""
+
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Iterator
+
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+from synkro.types.core import Trace
+
+console = Console()
+
+
+class Dataset(BaseModel):
+    """
+    A collection of generated training traces.
+
+    Provides methods for filtering, saving, and exporting traces
+    in various formats.
+
+    Examples:
+        >>> dataset = generator.generate(policy, traces=100)
+
+        >>> # Filter to only passing traces
+        >>> passing = dataset.filter(passed=True)
+
+        >>> # Save to JSONL
+        >>> dataset.save("training.jsonl")
+
+        >>> # Push to HuggingFace
+        >>> dataset.to_huggingface().push_to_hub("my-org/dataset")
+    """
+
+    traces: list[Trace] = Field(default_factory=list)
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __len__(self) -> int:
+        return len(self.traces)
+
+    def __iter__(self) -> Iterator[Trace]:
+        return iter(self.traces)
+
+    def __getitem__(self, idx: int) -> Trace:
+        return self.traces[idx]
+
+    def filter(
+        self,
+        passed: bool | None = None,
+        category: str | None = None,
+        min_length: int | None = None,
+    ) -> "Dataset":
+        """
+        Filter traces by criteria.
+
+        Args:
+            passed: Filter by grade pass/fail status
+            category: Filter by scenario category
+            min_length: Minimum response length in characters
+
+        Returns:
+            New Dataset with filtered traces
+        """
+        filtered = self.traces
+
+        if passed is not None:
+            filtered = [
+                t for t in filtered if t.grade and t.grade.passed == passed
+            ]
+
+        if category is not None:
+            filtered = [
+                t for t in filtered if t.scenario.category == category
+            ]
+
+        if min_length is not None:
+            filtered = [
+                t for t in filtered if len(t.assistant_message) >= min_length
+            ]
+
+        return Dataset(traces=filtered)
+
+    def dedupe(
+        self,
+        threshold: float = 0.85,
+        method: str = "semantic",
+        field: str = "user",
+    ) -> "Dataset":
+        """
+        Remove duplicate or near-duplicate traces.
+
+        Args:
+            threshold: Similarity threshold (0-1). Higher = stricter dedup.
+                       Only used for semantic method. (default: 0.85)
+            method: Deduplication method:
+                    - "exact": Remove exact text duplicates (fast)
+                    - "semantic": Remove semantically similar traces (requires sentence-transformers)
+            field: Which field to dedupe on - "user", "assistant", or "both"
+
+        Returns:
+            New Dataset with duplicates removed
+
+        Examples:
+            >>> # Remove exact duplicates (fast)
+            >>> deduped = dataset.dedupe(method="exact")
+
+            >>> # Remove semantically similar (needs sentence-transformers)
+            >>> deduped = dataset.dedupe(threshold=0.9, method="semantic")
+
+            >>> # Dedupe based on assistant responses
+            >>> deduped = dataset.dedupe(field="assistant")
+        """
+        if not self.traces:
+            return Dataset(traces=[])
+
+        if method == "exact":
+            return self._dedupe_exact(field)
+        elif method == "semantic":
+            return self._dedupe_semantic(threshold, field)
+        else:
+            raise ValueError(f"Unknown method: {method}. Use 'exact' or 'semantic'")
+
+    def _dedupe_exact(self, field: str) -> "Dataset":
+        """Remove exact text duplicates."""
+        seen = set()
+        unique_traces = []
+
+        for trace in self.traces:
+            if field == "user":
+                key = trace.user_message
+            elif field == "assistant":
+                key = trace.assistant_message
+            else:  # both
+                key = (trace.user_message, trace.assistant_message)
+
+            if key not in seen:
+                seen.add(key)
+                unique_traces.append(trace)
+
+        removed = len(self.traces) - len(unique_traces)
+        if removed > 0:
+            console.print(f"[yellow]🔍 Dedupe:[/yellow] Removed {removed} exact duplicates")
+
+        return Dataset(traces=unique_traces)
+
+    def _dedupe_semantic(self, threshold: float, field: str) -> "Dataset":
+        """Remove semantically similar traces using embeddings."""
+        try:
+            from sentence_transformers import SentenceTransformer
+            import numpy as np
+        except ImportError:
+            raise ImportError(
+                "sentence-transformers is required for semantic deduplication. "
+                "Install with: pip install sentence-transformers"
+            )
+
+        # Get texts to embed
+        if field == "user":
+            texts = [t.user_message for t in self.traces]
+        elif field == "assistant":
+            texts = [t.assistant_message for t in self.traces]
+        else:  # both
+            texts = [f"{t.user_message} {t.assistant_message}" for t in self.traces]
+
+        # Compute embeddings
+        console.print("[dim]Computing embeddings for deduplication...[/dim]")
+        model = SentenceTransformer("all-MiniLM-L6-v2")
+        embeddings = model.encode(texts, show_progress_bar=False)
+        embeddings = np.array(embeddings)
+
+        # Normalize for cosine similarity
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        embeddings = embeddings / norms
+
+        # Find duplicates using cosine similarity
+        unique_indices = []
+        duplicate_of = {}  # Maps duplicate index to original index
+
+        for i in range(len(embeddings)):
+            is_duplicate = False
+            for j in unique_indices:
+                similarity = np.dot(embeddings[i], embeddings[j])
+                if similarity >= threshold:
+                    is_duplicate = True
+                    duplicate_of[i] = j
+                    break
+            if not is_duplicate:
+                unique_indices.append(i)
+
+        unique_traces = [self.traces[i] for i in unique_indices]
+        removed = len(self.traces) - len(unique_traces)
+
+        if removed > 0:
+            console.print(f"[yellow]🔍 Dedupe:[/yellow] Removed {removed} semantic duplicates (threshold={threshold})")
+
+        return Dataset(traces=unique_traces)
+
+    @property
+    def passing_rate(self) -> float:
+        """Get the percentage of traces that passed grading."""
+        if not self.traces:
+            return 0.0
+
+        passed = sum(1 for t in self.traces if t.grade and t.grade.passed)
+        return passed / len(self.traces)
+
+    @property
+    def categories(self) -> list[str]:
+        """Get unique categories in the dataset."""
+        return list(set(t.scenario.category for t in self.traces if t.scenario.category))
+
+    def save(self, path: str | Path | None = None, format: str = "sft") -> "Dataset":
+        """
+        Save dataset to a JSONL file.
+
+        Args:
+            path: Output file path (auto-generated if not provided)
+            format: Output format - "sft", "qa", or "tool_call"
+
+        Returns:
+            Self for method chaining
+
+        Example:
+            >>> dataset.save()  # Auto-names: synkro_sft_2024-01-15.jsonl
+            >>> dataset.save("training.jsonl")
+            >>> dataset.save("qa_data.jsonl", format="qa")
+            >>> dataset.save("tools.jsonl", format="tool_call")
+        """
+        from synkro.formatters import SFTFormatter, QAFormatter, ToolCallFormatter
+
+        # Auto-generate filename if not provided
+        if path is None:
+            timestamp = datetime.now().strftime("%Y-%m-%d_%H%M")
+            path = f"synkro_{format}_{timestamp}.jsonl"
+        
+        path = Path(path)
+
+        if format == "sft":
+            SFTFormatter().save(self.traces, path)
+        elif format == "qa":
+            QAFormatter().save(self.traces, path)
+        elif format == "tool_call":
+            ToolCallFormatter().save(self.traces, path)
+        else:
+            raise ValueError(f"Unknown format: {format}. Use 'sft', 'qa', or 'tool_call'")
+        
+        # Print confirmation
+        file_size = path.stat().st_size
+        size_str = f"{file_size / 1024:.1f} KB" if file_size < 1024 * 1024 else f"{file_size / 1024 / 1024:.1f} MB"
+        console.print(f"[green]📁 Saved:[/green] {path} ({size_str})")
+        
+        return self
+
+    def to_jsonl(self, format: str = "sft") -> str:
+        """
+        Convert dataset to JSONL string.
+
+        Args:
+            format: Output format - "sft", "qa", or "tool_call"
+
+        Returns:
+            JSONL formatted string
+        """
+        from synkro.formatters import SFTFormatter, QAFormatter, ToolCallFormatter
+
+        if format == "sft":
+            return SFTFormatter().to_jsonl(self.traces)
+        elif format == "qa":
+            return QAFormatter().to_jsonl(self.traces)
+        elif format == "tool_call":
+            return ToolCallFormatter().to_jsonl(self.traces)
+        else:
+            raise ValueError(f"Unknown format: {format}. Use 'sft', 'qa', or 'tool_call'")
+
+    def to_hf_dataset(self, format: str = "sft"):
+        """
+        Convert to HuggingFace Dataset.
+
+        Args:
+            format: Output format - "sft", "qa", or "tool_call"
+
+        Returns:
+            HuggingFace datasets.Dataset object
+
+        Example:
+            >>> hf_dataset = dataset.to_hf_dataset()
+            >>> hf_dataset.push_to_hub("my-org/policy-traces")
+
+            >>> # With train/test split
+            >>> hf_dataset = dataset.to_hf_dataset()
+            >>> split = hf_dataset.train_test_split(test_size=0.1)
+            >>> split.push_to_hub("my-org/policy-traces")
+        """
+        try:
+            from datasets import Dataset as HFDataset
+        except ImportError:
+            raise ImportError(
+                "datasets is required for HuggingFace export. "
+                "Install with: pip install datasets"
+            )
+
+        from synkro.formatters import SFTFormatter, QAFormatter, ToolCallFormatter
+
+        if format == "sft":
+            examples = SFTFormatter(include_metadata=True).format(self.traces)
+        elif format == "qa":
+            examples = QAFormatter().format(self.traces)
+        elif format == "tool_call":
+            examples = ToolCallFormatter().format(self.traces)
+        else:
+            raise ValueError(f"Unknown format: {format}")
+
+        return HFDataset.from_list(examples)
+
+    # Alias for backwards compatibility
+    to_huggingface = to_hf_dataset
+
+    def push_to_hub(
+        self,
+        repo_id: str,
+        format: str = "sft",
+        private: bool = False,
+        split: str = "train",
+        token: str | None = None,
+    ) -> str:
+        """
+        Push dataset directly to HuggingFace Hub.
+
+        Args:
+            repo_id: HuggingFace repo ID (e.g., "my-org/policy-sft")
+            format: Output format - "sft", "qa", or "tool_call"
+            private: Whether the repo should be private
+            split: Dataset split name (default: "train")
+            token: HuggingFace token (uses cached token if not provided)
+
+        Returns:
+            URL of the uploaded dataset
+
+        Example:
+            >>> dataset.push_to_hub("my-org/policy-sft")
+            >>> dataset.push_to_hub("my-org/policy-sft", private=True)
+        """
+        hf_dataset = self.to_hf_dataset(format=format)
+        hf_dataset.push_to_hub(
+            repo_id,
+            private=private,
+            split=split,
+            token=token,
+        )
+        url = f"https://huggingface.co/datasets/{repo_id}"
+        console.print(f"[green]🤗 Pushed to Hub:[/green] {url}")
+        return url
+
+    def to_dict(self) -> dict:
+        """
+        Convert dataset to a dictionary.
+
+        Returns:
+            Dictionary with trace data
+        """
+        return {
+            "traces": [t.model_dump() for t in self.traces],
+            "stats": {
+                "total": len(self.traces),
+                "passing_rate": self.passing_rate,
+                "categories": self.categories,
+            },
+        }
+
+    def summary(self) -> str:
+        """
+        Get a summary of the dataset.
+
+        Returns:
+            Human-readable summary string
+        """
+        lines = [
+            f"Dataset Summary",
+            f"===============",
+            f"Total traces: {len(self.traces)}",
+            f"Passing rate: {self.passing_rate:.1%}",
+            f"Categories: {len(self.categories)}",
+        ]
+
+        if self.categories:
+            lines.append("")
+            lines.append("By category:")
+            for cat in self.categories:
+                count = sum(1 for t in self.traces if t.scenario.category == cat)
+                lines.append(f"  - {cat}: {count}")
+
+        return "\n".join(lines)
+
+    def __str__(self) -> str:
+        return f"Dataset(traces={len(self.traces)}, passing={self.passing_rate:.1%})"
+
+    def __repr__(self) -> str:
+        return self.__str__()
+