disco/pipeline/steps/step5_random_sampling.py

"""
Step 5: Random stratified sampling for attorney labeling.
"""

import random
from typing import List
import numpy as np
from pipeline.models.base import PipelineStep
from pipeline.common_defs import Chunk

class RandomSampler(PipelineStep):
    """Random stratified sampling for attorney labeling"""
    
    def __init__(self, n_samples: int = 20, seed: int = 42,
                 output_dir: str = './pipeline_output'):
        super().__init__(output_dir)
        self.n_samples = n_samples
        self.seed = seed
    
    def execute(self, chunks: List[Chunk]) -> List[Chunk]:
        """
        Select random stratified samples.
        
        Args:
            chunks: List of semantically filtered chunks
            
        Returns:
            List of sampled chunks
        """
        self.logger.info(f"Selecting {self.n_samples} random samples...")
        self.logger.info(f"Random seed: {self.seed}")
        
        random.seed(self.seed)
        
        # Stratify by semantic score quartiles
        scores = [c.semantic_score_combined for c in chunks if c.semantic_score_combined]
        
        if not scores:
            self.logger.warning("No semantic scores found, using random sampling")
            samples = random.sample(chunks, min(self.n_samples, len(chunks)))
        else:
            quartiles = np.percentile(scores, [25, 50, 75])
            samples = self._stratified_sample(chunks, quartiles)
        
        self.logger.info(f"Selected {len(samples)} samples")
        
        # Save samples
        self._save_samples(samples)
        
        return samples
    
    def _stratified_sample(self, chunks: List[Chunk], 
                          quartiles: np.ndarray) -> List[Chunk]:
        """Perform stratified sampling by score quartiles"""
        samples = []
        
        # Sample from each quartile
        for q_low, q_high in [(0, quartiles[0]), (quartiles[0], quartiles[1]),
                              (quartiles[1], quartiles[2]), (quartiles[2], 1.0)]:
            stratum = [
                c for c in chunks 
                if c.semantic_score_combined and 
                q_low <= c.semantic_score_combined < q_high
            ]
            
            if stratum:
                n_select = min(self.n_samples // 4, len(stratum))
                samples.extend(random.sample(stratum, n_select))
        
        # Fill remaining if needed
        if len(samples) < self.n_samples:
            remaining = [c for c in chunks if c not in samples]
            if remaining:
                n_more = min(self.n_samples - len(samples), len(remaining))
                samples.extend(random.sample(remaining, n_more))
        
        # Shuffle and limit
        random.shuffle(samples)
        return samples[:self.n_samples]
    
    def _save_samples(self, samples: List[Chunk]):
        """Save sampled chunks"""
        samples_data = [
            {
                'chunk_id': c.chunk_id,
                'start_line': c.start_line,
                'end_line': c.end_line,
                'semantic_score': c.semantic_score_combined,
                'num_messages': len(c.messages)
            }
            for c in samples
        ]
        
        self.save_results(samples_data, 'random_samples.json')

if __name__ == "__main__":
    # Example usage
    import json
    
    with open('pipeline_output/semantic_filtered_chunks.json', 'r') as f:
        data = json.load(f)
    
    # Reconstruct chunks (simplified for example)
    from pipeline.common_defs import Chunk, Message
    chunks = []
    for item in data['filtered_chunks']:
        chunk = Chunk(
            chunk_id=item['chunk_id'],
            start_line=item['start_line'],
            end_line=item['end_line'],
            messages=[],
            combined_text="",
            timestamp_start="",
            timestamp_end="",
            semantic_score_combined=item['score_combined']
        )
        chunks.append(chunk)
    
    sampler = RandomSampler(n_samples=20)
    samples = sampler.execute(chunks)
    print(f"Selected {len(samples)} samples")