real_dataset_loader.py

"""
Real Dataset Loader for Full MMLU, ARC, HellaSwag Evaluation
===========================================================

Loads actual benchmark datasets from HuggingFace instead of hardcoded samples.

Author: Principal Neuro-AI Engineer
Date: January 8, 2026
"""

import os
import json
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path
import logging

try:
    from datasets import load_dataset
    HF_DATASETS_AVAILABLE = True
except ImportError:
    HF_DATASETS_AVAILABLE = False
    print("[WARN]  HuggingFace datasets not installed. Install with: pip install datasets")

logger = logging.getLogger(__name__)

class RealDatasetLoader:
    """Load real benchmark datasets from HuggingFace or local cache"""
    
    def __init__(self, cache_dir: str = "./hf_cache"):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)
        
    def load_mmlu_full(self, split: str = "test", max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Load the full MMLU dataset (57 subjects)"""
        
        if not HF_DATASETS_AVAILABLE:
            return self._load_fallback_mmlu(max_samples)
        
        print(f"[DATA] Loading MMLU dataset (split: {split})...")
        
        try:
            # Load all MMLU subjects
            dataset = load_dataset("cais/mmlu", "all", cache_dir=self.cache_dir)
            
            samples = []
            subject_counts = {}
            
            for item in dataset[split]:
                subject = item['subject']
                if subject not in subject_counts:
                    subject_counts[subject] = 0
                
                sample = {
                    'question': item['question'],
                    'choices': item['choices'],
                    'answer': item['answer'],  # 0-3 index
                    'subject': subject,
                    'source': 'mmlu_full'
                }
                samples.append(sample)
                subject_counts[subject] += 1
                
                if max_samples and len(samples) >= max_samples:
                    break
            
            print(f"[OK] Loaded {len(samples)} MMLU samples from {len(subject_counts)} subjects")
            print(f"[CHART] Subjects: {list(subject_counts.keys())}")
            return samples
            
        except Exception as e:
            logger.error(f"Failed to load MMLU: {e}")
            print(f"[FAIL] Failed to load MMLU dataset: {e}")
            return self._load_fallback_mmlu(max_samples)
    
    def load_mmlu_subject(self, subject: str, split: str = "test", max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Load MMLU questions for a specific subject"""
        
        if not HF_DATASETS_AVAILABLE:
            return self._load_fallback_mmlu_subject(subject, max_samples)
        
        print(f"[DATA] Loading MMLU {subject} (split: {split})...")
        
        try:
            # Load specific subject
            dataset = load_dataset("cais/mmlu", subject, cache_dir=self.cache_dir)
            
            samples = []
            for item in dataset[split]:
                sample = {
                    'question': item['question'],
                    'choices': item['choices'],
                    'answer': item['answer'],  # 0-3 index
                    'subject': subject,
                    'source': f'mmlu_{subject}'
                }
                samples.append(sample)
                
                if max_samples and len(samples) >= max_samples:
                    break
            
            print(f"[OK] Loaded {len(samples)} {subject} samples")
            return samples
            
        except Exception as e:
            logger.error(f"Failed to load MMLU {subject}: {e}")
            print(f"[FAIL] Failed to load MMLU {subject}: {e}")
            return self._load_fallback_mmlu_subject(subject, max_samples)
    
    def load_arc_full(self, split: str = "test", max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Load the full ARC (AI2 Reasoning Challenge) dataset"""
        
        if not HF_DATASETS_AVAILABLE:
            return self._load_fallback_arc(max_samples)
            
        print(f"[PUZZLE] Loading ARC dataset (split: {split})...")
        
        try:
            # Load ARC-Challenge (harder version)
            dataset = load_dataset("ai2_arc", "ARC-Challenge", cache_dir=self.cache_dir)
            
            samples = []
            for item in dataset[split]:
                # Convert choices to list format
                choices_data = item['choices']
                if isinstance(choices_data, dict):
                    choices = choices_data['text']
                    labels = choices_data['label']
                else:
                    choices = choices_data
                    labels = ['A', 'B', 'C', 'D'][:len(choices)]
                
                # Find correct answer index
                correct_label = item['answerKey']
                try:
                    answer_idx = labels.index(correct_label)
                except (ValueError, TypeError):
                    # Handle cases where answerKey is 1,2,3,4 instead of A,B,C,D
                    label_map = {'A': 0, 'B': 1, 'C': 2, 'D': 3, '1': 0, '2': 1, '3': 2, '4': 3}
                    answer_idx = label_map.get(correct_label, 0)
                
                sample = {
                    'question': item['question'],
                    'choices': choices,
                    'answer': answer_idx,
                    'subject': 'science_reasoning',
                    'source': 'arc_challenge',
                    'id': item.get('id', f"arc_{len(samples)}")
                }
                samples.append(sample)
                
                if max_samples and len(samples) >= max_samples:
                    break
                    
            print(f"[OK] Loaded {len(samples)} ARC-Challenge samples")
            return samples
            
        except Exception as e:
            logger.error(f"Failed to load ARC: {e}")
            print(f"[FAIL] Failed to load ARC dataset: {e}")
            return self._load_fallback_arc(max_samples)
    
    def load_hellaswag_full(self, split: str = "validation", max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Load the full HellaSwag dataset"""
        
        if not HF_DATASETS_AVAILABLE:
            return self._load_fallback_hellaswag(max_samples)
            
        print(f"[THINK] Loading HellaSwag dataset (split: {split})...")
        
        try:
            dataset = load_dataset("hellaswag", cache_dir=self.cache_dir)
            
            samples = []
            for item in dataset[split]:
                # HellaSwag format: ctx + endings
                context = item['ctx']
                endings = item['endings']
                
                # Format as multiple choice
                question = f"Context: {context}\n\nWhat happens next?"
                choices = [f"{i+1}. {ending}" for i, ending in enumerate(endings)]
                
                sample = {
                    'question': question,
                    'choices': choices,
                    'answer': int(item['label']),  # Correct ending index
                    'subject': 'commonsense_reasoning',
                    'source': 'hellaswag',
                    'context': context,
                    'endings': endings
                }
                samples.append(sample)
                
                if max_samples and len(samples) >= max_samples:
                    break
                    
            print(f"[OK] Loaded {len(samples)} HellaSwag samples")
            return samples
            
        except Exception as e:
            logger.error(f"Failed to load HellaSwag: {e}")
            print(f"[FAIL] Failed to load HellaSwag dataset: {e}")
            return self._load_fallback_hellaswag(max_samples)
    
    def _load_fallback_mmlu(self, max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Fallback MMLU samples if HF datasets unavailable"""
        samples = [
            {
                'question': 'What is the derivative of ln(x)?',
                'choices': ['1/x', 'ln(x)', 'x', 'e^x'],
                'answer': 0,
                'subject': 'mathematics',
                'source': 'fallback'
            },
            {
                'question': 'Which organelle is responsible for cellular respiration?',
                'choices': ['Nucleus', 'Mitochondria', 'Ribosome', 'Golgi apparatus'],
                'answer': 1,
                'subject': 'biology',
                'source': 'fallback'
            },
            {
                'question': 'What is the time complexity of binary search?',
                'choices': ['O(n)', 'O(log n)', 'O(n²)', 'O(1)'],
                'answer': 1,
                'subject': 'computer_science',
                'source': 'fallback'
            }
        ]
        
        if max_samples:
            samples = samples[:max_samples]
            
        print(f"[WARN]  Using fallback MMLU samples: {len(samples)} questions")
        return samples
    
    def _load_fallback_arc(self, max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Fallback ARC samples if HF datasets unavailable"""
        samples = [
            {
                'question': 'A student wants to know which part of a plant cell stores water. Which part should the student study?',
                'choices': ['cell wall', 'vacuole', 'nucleus', 'chloroplast'],
                'answer': 1,
                'subject': 'science_reasoning',
                'source': 'fallback'
            },
            {
                'question': 'Which process is responsible for changing liquid water into water vapor?',
                'choices': ['condensation', 'evaporation', 'precipitation', 'freezing'],
                'answer': 1,
                'subject': 'science_reasoning',
                'source': 'fallback'
            }
        ]
        
        if max_samples:
            samples = samples[:max_samples]
            
        print(f"[WARN]  Using fallback ARC samples: {len(samples)} questions")
        return samples
    
    def _load_fallback_hellaswag(self, max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Fallback HellaSwag samples if HF datasets unavailable"""
        samples = [
            {
                'question': 'Context: A person is cooking pasta in a kitchen.\n\nWhat happens next?',
                'choices': [
                    '1. They drain the pasta and serve it with sauce.',
                    '2. They put the pasta in the refrigerator.',
                    '3. They throw the pasta in the garbage.',
                    '4. They use the pasta to clean the floor.'
                ],
                'answer': 0,
                'subject': 'commonsense_reasoning',
                'source': 'fallback'
            }
        ]
        
        if max_samples:
            samples = samples[:max_samples]
            
        print(f"[WARN]  Using fallback HellaSwag samples: {len(samples)} questions")
        return samples
    
    def _load_fallback_mmlu_subject(self, subject: str, max_samples: Optional[int] = None) -> List[Dict[str, Any]]:
        """Fallback abstract algebra samples when HuggingFace is not available"""
        
        # Abstract algebra questions (simplified versions)
        algebra_questions = [
            {
                'question': 'Let G be a group and H a subgroup of G. Which of the following is true about the index [G:H]?',
                'choices': ['It is always finite', 'It can be infinite', 'It is always 1', 'It depends on the subgroup'],
                'answer': 1,  # B
                'subject': 'abstract_algebra'
            },
            {
                'question': 'What is the order of the symmetric group S_3?',
                'choices': ['3', '6', '9', '12'],
                'answer': 1,  # B (6)
                'subject': 'abstract_algebra'
            },
            {
                'question': 'Which of the following is a field?',
                'choices': ['ℤ (integers)', 'ℚ (rationals)', 'ℤ/2ℤ', 'All of the above'],
                'answer': 3,  # D
                'subject': 'abstract_algebra'
            },
            {
                'question': 'If phi: G → H is a group homomorphism, then ker(phi) is:',
                'choices': ['A subgroup of G', 'A subgroup of H', 'The image of phi', 'None of the above'],
                'answer': 0,  # A
                'subject': 'abstract_algebra'
            },
            {
                'question': 'The center of a group G consists of elements that:',
                'choices': ['Commute with all elements', 'Are in the center of G', 'Generate G', 'Have order 2'],
                'answer': 0,  # A
                'subject': 'abstract_algebra'
            },
            {
                'question': 'A ring is a set with two binary operations that satisfy:',
                'choices': ['Group axioms', 'Field axioms', 'Ring axioms', 'Vector space axioms'],
                'answer': 2,  # C
                'subject': 'abstract_algebra'
            },
            {
                'question': 'The fundamental theorem of homomorphism for groups states that:',
                'choices': ['G/ker(phi) ≅ im(phi)', 'G ≅ H for any homomorphism', 'All homomorphisms are isomorphisms', 'Groups have unique homomorphisms'],
                'answer': 0,  # A
                'subject': 'abstract_algebra'
            },
            {
                'question': 'Which of the following groups is abelian?',
                'choices': ['S_3', 'ℤ_2 x ℤ_2', 'GL(2,ℝ)', 'A_4'],
                'answer': 1,  # B
                'subject': 'abstract_algebra'
            },
            {
                'question': 'The order of an element g in a group is:',
                'choices': ['The number of elements in the group', 'The smallest positive integer n such that g^n = e', 'The index of the subgroup generated by g', 'Always finite'],
                'answer': 1,  # B
                'subject': 'abstract_algebra'
            },
            {
                'question': 'A normal subgroup N of G satisfies:',
                'choices': ['gN = Ng for all g in G', 'N is a subgroup', 'G/N is a group', 'All of the above'],
                'answer': 3,  # D
                'subject': 'abstract_algebra'
            }
        ]
        
        samples = algebra_questions
        if max_samples:
            samples = samples[:max_samples]
            
        print(f"[WARN]  Using fallback {subject} samples: {len(samples)} questions")
        return samples
    
    def load_all_benchmarks(self, max_samples_per_benchmark: Optional[int] = None) -> Dict[str, List[Dict[str, Any]]]:
        """Load all benchmark datasets"""
        
        print("[ROCKET] Loading all benchmark datasets...")
        
        benchmarks = {
            'mmlu': self.load_mmlu_full(max_samples=max_samples_per_benchmark),
            'arc': self.load_arc_full(max_samples=max_samples_per_benchmark),
            'hellaswag': self.load_hellaswag_full(max_samples=max_samples_per_benchmark)
        }
        
        total_samples = sum(len(samples) for samples in benchmarks.values())
        print(f"[OK] Total loaded: {total_samples} samples across {len(benchmarks)} benchmarks")
        
        return benchmarks

def install_datasets_if_needed():
    """Install required packages if not available"""
    if not HF_DATASETS_AVAILABLE:
        print("? Installing required packages...")
        import subprocess
        import sys
        
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", "datasets", "huggingface_hub"])
            print("[OK] Packages installed successfully!")
            print("[CYCLE] Please restart the script to use HuggingFace datasets")
            return True
        except subprocess.CalledProcessError as e:
            print(f"[FAIL] Failed to install packages: {e}")
            return False
    return True

if __name__ == "__main__":
    # Test the loader
    if not install_datasets_if_needed():
        exit(1)
    
    loader = RealDatasetLoader()
    
    # Test loading each dataset
    print("\n" + "="*60)
    print("Testing Real Dataset Loader")
    print("="*60)
    
    mmlu_samples = loader.load_mmlu_full(max_samples=5)
    print(f"\nMMlu sample: {mmlu_samples[0] if mmlu_samples else 'None'}")
    
    arc_samples = loader.load_arc_full(max_samples=5)
    print(f"\nARC sample: {arc_samples[0] if arc_samples else 'None'}")
    
    hellaswag_samples = loader.load_hellaswag_full(max_samples=5)
    print(f"\nHellaSwag sample: {hellaswag_samples[0] if hellaswag_samples else 'None'}")