gradia 1.0.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

gradia/events/tracker.py

@@ -0,0 +1,337 @@
+"""
+Sample Tracker for Learning Timeline (v2.0.0)
+
+Manages which samples to track and maintains their state across epochs.
+Implements deterministic, bounded sampling strategy.
+"""
+
+from typing import List, Dict, Any, Optional, Set
+import numpy as np
+from dataclasses import dataclass, field
+
+from .models import LearningEvent, SampleState, EpochSummary
+
+
+@dataclass
+class SampleTracker:
+    """
+    Tracks a bounded subset of samples throughout training.
+    
+    Strategy:
+    - Auto-select hard samples (near decision boundary)
+    - Include user-selected samples if specified
+    - Cap at max_samples for performance
+    - Use deterministic seeding for reproducibility
+    
+    Attributes:
+        max_samples: Maximum samples to track (default 100)
+        seed: Random seed for reproducibility
+        tracked_indices: Set of sample indices being tracked
+        sample_states: State history for each tracked sample
+    """
+    max_samples: int = 100
+    seed: int = 42
+    tracked_indices: Set[int] = field(default_factory=set)
+    sample_states: Dict[int, SampleState] = field(default_factory=dict)
+    user_selected: Set[int] = field(default_factory=set)
+    run_id: str = ""
+    _initialized: bool = False
+    
+    def initialize(
+        self,
+        X: np.ndarray,
+        y: np.ndarray,
+        run_id: str,
+        user_indices: Optional[List[int]] = None,
+        model: Optional[Any] = None
+    ):
+        """
+        Initialize tracking with dataset and optional model predictions.
+        
+        Args:
+            X: Feature matrix
+            y: Labels
+            run_id: Unique run identifier
+            user_indices: User-selected sample indices to always track
+            model: Optional model for boundary sample selection
+        """
+        self.run_id = run_id
+        self._rng = np.random.RandomState(self.seed)
+        
+        n_samples = len(y)
+        
+        # Start with user-selected samples
+        if user_indices:
+            self.user_selected = set(user_indices[:self.max_samples // 2])
+            self.tracked_indices = self.user_selected.copy()
+        
+        remaining_slots = self.max_samples - len(self.tracked_indices)
+        
+        if remaining_slots > 0:
+            # Try to select boundary/hard samples if model available
+            if model is not None and hasattr(model, 'predict_proba'):
+                boundary_indices = self._select_boundary_samples(X, y, model, remaining_slots)
+                self.tracked_indices.update(boundary_indices)
+            
+            # Fill remaining with stratified random
+            remaining_slots = self.max_samples - len(self.tracked_indices)
+            if remaining_slots > 0:
+                available = set(range(n_samples)) - self.tracked_indices
+                random_indices = self._stratified_sample(
+                    list(available), y, remaining_slots
+                )
+                self.tracked_indices.update(random_indices)
+        
+        # Initialize sample states
+        for idx in self.tracked_indices:
+            self.sample_states[idx] = SampleState(
+                sample_id=idx,
+                true_label=y[idx] if hasattr(y, '__getitem__') else y.iloc[idx]
+            )
+        
+        self._initialized = True
+        
+    def _select_boundary_samples(
+        self,
+        X: np.ndarray,
+        y: np.ndarray,
+        model: Any,
+        n_samples: int
+    ) -> Set[int]:
+        """
+        Select samples near the decision boundary.
+        
+        These are the most informative for understanding model learning.
+        """
+        try:
+            probas = model.predict_proba(X)
+            
+            # Compute margin: difference between top 2 class probabilities
+            if probas.shape[1] >= 2:
+                sorted_probas = np.sort(probas, axis=1)
+                margins = sorted_probas[:, -1] - sorted_probas[:, -2]
+            else:
+                margins = np.abs(probas[:, 0] - 0.5)
+            
+            # Lower margin = closer to boundary = more interesting
+            # Exclude already tracked
+            available_mask = np.ones(len(margins), dtype=bool)
+            for idx in self.tracked_indices:
+                available_mask[idx] = False
+            
+            margins[~available_mask] = np.inf
+            
+            # Select lowest margin samples
+            boundary_indices = np.argsort(margins)[:n_samples]
+            return set(boundary_indices.tolist())
+            
+        except Exception:
+            # Fallback if predict_proba fails
+            return set()
+    
+    def _stratified_sample(
+        self,
+        available: List[int],
+        y: np.ndarray,
+        n_samples: int
+    ) -> Set[int]:
+        """
+        Stratified random sampling to maintain class balance.
+        """
+        if not available:
+            return set()
+            
+        # Group by class
+        class_indices: Dict[Any, List[int]] = {}
+        for idx in available:
+            label = y[idx] if hasattr(y, '__getitem__') else y.iloc[idx]
+            if label not in class_indices:
+                class_indices[label] = []
+            class_indices[label].append(idx)
+        
+        # Sample proportionally from each class
+        selected = []
+        n_classes = len(class_indices)
+        per_class = max(1, n_samples // n_classes)
+        
+        for label, indices in class_indices.items():
+            k = min(per_class, len(indices))
+            sampled = self._rng.choice(indices, size=k, replace=False)
+            selected.extend(sampled.tolist())
+        
+        # Trim if over budget
+        if len(selected) > n_samples:
+            selected = self._rng.choice(selected, size=n_samples, replace=False).tolist()
+            
+        return set(selected)
+    
+    def record_predictions(
+        self,
+        epoch: int,
+        X: np.ndarray,
+        y: np.ndarray,
+        predictions: np.ndarray,
+        probabilities: Optional[np.ndarray] = None
+    ) -> List[LearningEvent]:
+        """
+        Record predictions for all tracked samples at this epoch.
+        
+        Args:
+            epoch: Current epoch number
+            X: Full feature matrix
+            y: Full labels
+            predictions: Model predictions for all samples
+            probabilities: Optional probability matrix
+            
+        Returns:
+            List of LearningEvents for this epoch
+        """
+        if not self._initialized:
+            raise RuntimeError("SampleTracker not initialized. Call initialize() first.")
+        
+        events = []
+        
+        for idx in self.tracked_indices:
+            true_label = y[idx] if hasattr(y, '__getitem__') else y.iloc[idx]
+            pred_label = predictions[idx]
+            
+            # Compute confidence
+            if probabilities is not None:
+                proba_row = probabilities[idx]
+                confidence = float(np.max(proba_row))
+                proba_list = proba_row.tolist()
+                
+                # Compute margin
+                sorted_p = np.sort(proba_row)
+                margin = float(sorted_p[-1] - sorted_p[-2]) if len(sorted_p) >= 2 else confidence
+            else:
+                confidence = 1.0  # No probability info
+                proba_list = None
+                margin = None
+            
+            correct = (pred_label == true_label)
+            
+            event = LearningEvent(
+                run_id=self.run_id,
+                epoch=epoch,
+                sample_id=idx,
+                true_label=true_label,
+                predicted_label=pred_label,
+                confidence=confidence,
+                correct=bool(correct),
+                margin=margin,
+                probabilities=proba_list
+            )
+            
+            # Update sample state
+            self.sample_states[idx].add_event(event)
+            events.append(event)
+        
+        return events
+    
+    def get_epoch_summary(self, epoch: int) -> EpochSummary:
+        """
+        Generate aggregated summary for an epoch.
+        """
+        states = list(self.sample_states.values())
+        
+        # Count by stability class
+        stability_counts = {
+            "stable_correct": 0,
+            "stable_wrong": 0,
+            "unstable": 0,
+            "late_learner": 0,
+            "unknown": 0
+        }
+        
+        correct_count = 0
+        flip_count = 0
+        
+        for state in states:
+            stability_counts[state.stability_class] += 1
+            if state.history and state.history[-1].correct:
+                correct_count += 1
+            flip_count += state.flip_count
+        
+        return EpochSummary(
+            run_id=self.run_id,
+            epoch=epoch,
+            timestamp=__import__('time').time(),
+            total_tracked=len(states),
+            correct_count=correct_count,
+            flip_count=flip_count,
+            stable_correct=stability_counts["stable_correct"],
+            stable_wrong=stability_counts["stable_wrong"],
+            unstable=stability_counts["unstable"],
+            late_learners=stability_counts["late_learner"]
+        )
+    
+    def get_top_flipping_samples(self, n: int = 10) -> List[SampleState]:
+        """Get samples with most prediction flips."""
+        sorted_states = sorted(
+            self.sample_states.values(),
+            key=lambda s: s.flip_count,
+            reverse=True
+        )
+        return sorted_states[:n]
+    
+    def get_late_learners(self, threshold_epoch: int = 5) -> List[SampleState]:
+        """Get samples that became correct after threshold epoch."""
+        late = []
+        for state in self.sample_states.values():
+            first = state.first_correct_epoch
+            if first is not None and first >= threshold_epoch:
+                late.append(state)
+        return sorted(late, key=lambda s: s.first_correct_epoch or 999)
+    
+    def get_never_correct(self) -> List[SampleState]:
+        """Get samples that were never correctly classified."""
+        return [
+            state for state in self.sample_states.values()
+            if state.first_correct_epoch is None and state.history
+        ]
+    
+    def get_sample_state(self, sample_id: int) -> Optional[SampleState]:
+        """Get state for a specific sample."""
+        return self.sample_states.get(sample_id)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize tracker state for storage."""
+        return {
+            "max_samples": self.max_samples,
+            "seed": self.seed,
+            "run_id": self.run_id,
+            "tracked_indices": list(self.tracked_indices),
+            "user_selected": list(self.user_selected),
+            "sample_states": {
+                idx: {
+                    "sample_id": state.sample_id,
+                    "true_label": state.true_label,
+                    "history": [e.to_dict() for e in state.history]
+                }
+                for idx, state in self.sample_states.items()
+            }
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SampleTracker":
+        """Restore tracker from serialized state."""
+        tracker = cls(
+            max_samples=data["max_samples"],
+            seed=data["seed"],
+            run_id=data["run_id"]
+        )
+        tracker.tracked_indices = set(data["tracked_indices"])
+        tracker.user_selected = set(data.get("user_selected", []))
+        
+        for idx_str, state_data in data["sample_states"].items():
+            idx = int(idx_str)
+            state = SampleState(
+                sample_id=state_data["sample_id"],
+                true_label=state_data["true_label"],
+                history=[LearningEvent.from_dict(e) for e in state_data["history"]]
+            )
+            tracker.sample_states[idx] = state
+        
+        tracker._initialized = True
+        return tracker

gradia/trainer/engine.py

@@ -15,6 +15,7 @@ from ..models.base import GradiaModel
 from ..models.sklearn_wrappers import ModelFactory
 from ..core.scenario import Scenario
 from .callbacks import Callback, EventLogger
+from ..events import SampleTracker, TimelineLogger
 
 class Trainer:
     def __init__(self, scenario: Scenario, config: Dict[str, Any], run_dir: str):
@@ -28,6 +29,25 @@ class Trainer:
             config['model'].get('params', {})
         )
         self.callbacks: List[Callback] = [EventLogger(run_dir)]
+        
+        # v2.0: Learning Timeline support
+        self.enable_timeline = config.get('timeline', {}).get('enabled', True)
+        self.sample_tracker: SampleTracker = None
+        self.timeline_logger: TimelineLogger = None
+        
+        if self.enable_timeline:
+            timeline_config = config.get('timeline', {})
+            self.sample_tracker = SampleTracker(
+                max_samples=timeline_config.get('max_samples', 100),
+                seed=config['training'].get('random_seed', 42)
+            )
+            self.timeline_logger = TimelineLogger(run_dir)
+        
+        # Store data references for evaluation
+        self._X_test = None
+        self._y_test = None
+        self._X_train = None
+        self._y_train = None
 
     def run(self):
         print("DEBUG: Trainer.run() started.")
@@ -82,6 +102,25 @@ class Trainer:
                 random_state=self.config['training'].get('random_seed', 42)
             )
             
+            # Store for evaluation
+            self._X_train = X_train
+            self._y_train = y_train
+            self._X_test = X_test
+            self._y_test = y_test
+            
+            # v2.0: Initialize sample tracker for Learning Timeline
+            run_id = f"run_{int(time.time())}"
+            if self.enable_timeline and self.sample_tracker:
+                user_samples = self.config.get('timeline', {}).get('user_samples', None)
+                self.sample_tracker.initialize(
+                    X=X_test.values if hasattr(X_test, 'values') else X_test,
+                    y=y_test.values if hasattr(y_test, 'values') else y_test,
+                    run_id=run_id,
+                    user_indices=user_samples
+                )
+                # Clear previous timeline logs for fresh run
+                self.timeline_logger.clear()
+            
             # Notify Start
             epochs = self.config['training'].get('epochs', 10)
             self._dispatch('on_train_begin', {
@@ -103,7 +142,6 @@ class Trainer:
                 classes = np.unique(y) if self.scenario.task_type == 'classification' else None
                 
                 # TQDM Output to Console
-                import time
                 with tqdm(range(1, epochs + 1), desc="Training", unit="epoch", colour="green") as pbar:
                     for epoch in pbar:
                         # Small delay to visualize speed if too fast
@@ -113,6 +151,10 @@ class Trainer:
                         
                         # Evaluate
                         metrics = self._evaluate(X_train, y_train, X_test, y_test)
+                        
+                        # v2.0: Record timeline events
+                        self._record_timeline_epoch(epoch, X_test, y_test)
+                        
                         self._dispatch('on_epoch_end', epoch, metrics)
     
             else:
@@ -125,12 +167,15 @@ class Trainer:
                 metrics = self._evaluate(X_train, y_train, X_test, y_test)
                 
                 # Simulate progress bar so UI doesn't look broken
-                import time
                 with tqdm(range(1, epochs + 1), desc="Training", unit="epoch", colour="blue") as pbar:
                     for epoch in pbar:
                         time.sleep(0.1) # Simulate work
                         # We broadcast the SAME metrics for every "epoch" since the model doesn't change
                         # But it keeps the UI happy and consistent
+                        
+                        # v2.0: Record timeline events (same predictions each epoch for non-iterative)
+                        self._record_timeline_epoch(epoch, X_test, y_test)
+                        
                         self._dispatch('on_epoch_end', epoch, metrics)
                         
                         # Update Progress Bar
@@ -145,11 +190,25 @@ class Trainer:
             
             # 5. Finalize
             fi = self.model.get_feature_importance()
+            
+            # v2.0: Finalize timeline logging
+            if self.enable_timeline and self.timeline_logger:
+                self.timeline_logger.save_tracker_state(self.sample_tracker.to_dict())
+                self.timeline_logger.finalize()
         
             # 6. Training Complete
+            timeline_summary = None
+            if self.enable_timeline and self.sample_tracker:
+                timeline_summary = {
+                    "tracked_samples": len(self.sample_tracker.tracked_indices),
+                    "top_flipping": [s.sample_id for s in self.sample_tracker.get_top_flipping_samples(5)],
+                    "never_correct": len(self.sample_tracker.get_never_correct())
+                }
+            
             self._dispatch("on_train_end", {
                 "epoch": epochs, 
-                "feature_importance": self.model.get_feature_importance()
+                "feature_importance": self.model.get_feature_importance(),
+                "timeline_summary": timeline_summary
             })
             
             # 7. Save Model
@@ -201,3 +260,116 @@ class Trainer:
     def _dispatch(self, method_name, *args, **kwargs):
         for cb in self.callbacks:
             getattr(cb, method_name)(*args, **kwargs)
+
+    def _record_timeline_epoch(self, epoch: int, X_test, y_test):
+        """
+        Record sample-level predictions for Learning Timeline (v2.0).
+        
+        Captures predictions for all tracked samples and logs events.
+        """
+        if not self.enable_timeline or not self.sample_tracker:
+            return
+        
+        try:
+            # Get predictions for all test samples
+            X_arr = X_test.values if hasattr(X_test, 'values') else X_test
+            y_arr = y_test.values if hasattr(y_test, 'values') else y_test
+            
+            predictions = self.model.predict(X_arr)
+            
+            # Get probabilities if available
+            probabilities = None
+            if hasattr(self.model, 'predict_proba'):
+                try:
+                    probabilities = self.model.predict_proba(X_arr)
+                except Exception:
+                    pass
+            
+            # Record predictions for tracked samples
+            events = self.sample_tracker.record_predictions(
+                epoch=epoch,
+                X=X_arr,
+                y=y_arr,
+                predictions=predictions,
+                probabilities=probabilities
+            )
+            
+            # Log events
+            self.timeline_logger.log_events(events, flush=(epoch % 5 == 0))
+            
+            # Log epoch summary
+            summary = self.sample_tracker.get_epoch_summary(epoch)
+            self.timeline_logger.log_summary(summary)
+            
+        except Exception as e:
+            print(f"Warning: Timeline recording failed for epoch {epoch}: {e}")
+
+    def evaluate_full(self) -> Dict[str, Any]:
+        """
+        Run full evaluation on test set with confusion matrix.
+        
+        Returns:
+            Dictionary with evaluation results including confusion matrix.
+        """
+        if self._X_test is None or self._y_test is None:
+            raise ValueError("No test data available. Run training first.")
+        
+        results = {}
+        
+        preds = self.model.predict(self._X_test)
+        
+        if self.scenario.task_type == 'classification':
+            results['accuracy'] = accuracy_score(self._y_test, preds)
+            results['precision'] = precision_score(self._y_test, preds, average='weighted', zero_division=0)
+            results['recall'] = recall_score(self._y_test, preds, average='weighted', zero_division=0)
+            results['f1'] = f1_score(self._y_test, preds, average='weighted', zero_division=0)
+            
+            # Confusion matrix
+            cm = confusion_matrix(self._y_test, preds)
+            results['confusion_matrix'] = cm.tolist()
+            
+            # Class labels
+            classes = sorted(list(set(self._y_test.tolist() if hasattr(self._y_test, 'tolist') else self._y_test)))
+            results['classes'] = [str(c) for c in classes]
+        else:
+            results['mse'] = mean_squared_error(self._y_test, preds)
+            results['mae'] = mean_absolute_error(self._y_test, preds)
+            results['r2'] = r2_score(self._y_test, preds)
+        
+        # v2.0: Include timeline insights if available
+        if self.enable_timeline and self.sample_tracker:
+            results['timeline_insights'] = {
+                'total_tracked': len(self.sample_tracker.tracked_indices),
+                'top_flipping': [
+                    {
+                        'sample_id': s.sample_id,
+                        'flip_count': s.flip_count,
+                        'true_label': str(s.true_label),
+                        'stability': s.stability_class
+                    }
+                    for s in self.sample_tracker.get_top_flipping_samples(10)
+                ],
+                'never_correct_count': len(self.sample_tracker.get_never_correct()),
+                'late_learners_count': len(self.sample_tracker.get_late_learners())
+            }
+        
+        return results
+
+    def get_sample_timeline(self, sample_id: int) -> List[Dict[str, Any]]:
+        """
+        Get the learning timeline for a specific sample.
+        
+        Args:
+            sample_id: Index of the sample to retrieve
+            
+        Returns:
+            List of event dictionaries showing prediction evolution
+        """
+        if not self.enable_timeline or not self.sample_tracker:
+            return []
+        
+        state = self.sample_tracker.get_sample_state(sample_id)
+        if not state:
+            return []
+        
+        return [event.to_dict() for event in state.history]