ToGMAL MCP Server

""" Enhanced Dataset Fetcher for ToGMAL Clustering Fetches datasets categorized into GOOD, LIMITATIONS, and HARMFUL clusters """ import asyncio import json import logging from pathlib import Path from typing import List, Dict, Any, Optional from dataclasses import dataclass, asdict from datetime import datetime logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Try to import datasets, fall back gracefully try: from datasets import load_dataset HAS_DATASETS = True except ImportError: logger.warning("datasets library not installed. Install with: uv pip install datasets") HAS_DATASETS = False @dataclass class DatasetConfig: """Configuration for a dataset source""" name: str source_id: str # HuggingFace dataset name split: str = "train" text_column: str = "text" max_samples: int = 1000 # Limit for faster iteration # Cluster classification cluster_category: str = "unknown" # "good", "limitations", "harmful" domain: str = "general" # Performance metrics (if known) llm_performance: Optional[float] = None # 0-1, e.g., 0.42 for 42% accuracy # ============================================================================ # Dataset Catalog - Organized by Expected Cluster # ============================================================================ DATASET_CATALOG = { # ======================================================================= # GOOD CLUSTER: High LLM performance (>80% accuracy) # ======================================================================= "good": [ DatasetConfig( name="squad_general_qa", source_id="rajpurkar/squad_v2", split="validation", text_column="question", cluster_category="good", domain="general_qa", llm_performance=0.86, max_samples=500 ), DatasetConfig( name="hellaswag_commonsense", source_id="Rowan/hellaswag", split="validation", text_column="ctx", cluster_category="good", domain="commonsense", llm_performance=0.95, max_samples=500 ), # Note: WMT14 and CNN/DailyMail are large, starting with smaller ones ], # ======================================================================= # LIMITATIONS CLUSTER: Poor LLM performance (<70% accuracy) # ======================================================================= "limitations": [ DatasetConfig( name="math_competition", source_id="hendrycks/competition_math", split="test", text_column="problem", cluster_category="limitations", domain="mathematics", llm_performance=0.42, max_samples=500 ), DatasetConfig( name="medical_qa", source_id="GBaker/MedQA-USMLE-4-options", split="test", text_column="question", cluster_category="limitations", domain="medicine", llm_performance=0.65, max_samples=500 ), DatasetConfig( name="code_defects", source_id="code_x_glue_cc_defect_detection", split="test", text_column="func", cluster_category="limitations", domain="coding", llm_performance=0.60, # Estimated max_samples=500 ), ], # ======================================================================= # HARMFUL CLUSTER: Safety benchmarks (jailbreaks, toxic content) # ======================================================================= "harmful": [ DatasetConfig( name="toxic_chat", source_id="lmsys/toxic-chat", split="train", text_column="user_input", cluster_category="harmful", domain="safety", llm_performance=None, # N/A for safety max_samples=500 ), # Note: hh-rlhf is large, will use smaller sample ], } @dataclass class DatasetEntry: """Single entry from a dataset""" id: str text: str cluster_category: str # "good", "limitations", "harmful" domain: str source: str metadata: Dict[str, Any] = None def __post_init__(self): if self.metadata is None: self.metadata = {} if not self.id: import hashlib self.id = hashlib.sha256(self.text.encode()).hexdigest()[:16] class EnhancedDatasetFetcher: """ Fetches datasets for clustering analysis Organizes into GOOD, LIMITATIONS, and HARMFUL categories """ def __init__(self, cache_dir: Path = Path("./data/datasets")): self.cache_dir = cache_dir self.cache_dir.mkdir(parents=True, exist_ok=True) logger.info(f"EnhancedDatasetFetcher initialized (cache: {cache_dir})") async def fetch_all_datasets(self) -> Dict[str, List[DatasetEntry]]: """ Fetch all datasets organized by cluster category Returns: { "good": [DatasetEntry, ...], "limitations": [DatasetEntry, ...], "harmful": [DatasetEntry, ...] } """ if not HAS_DATASETS: logger.error("datasets library not installed!") logger.info("Run: uv pip install datasets") return self._generate_synthetic_data() all_data = {"good": [], "limitations": [], "harmful": []} for category, configs in DATASET_CATALOG.items(): logger.info(f"\n{'='*60}") logger.info(f"Fetching {category.upper()} cluster datasets") logger.info(f"{'='*60}") for config in configs: try: entries = await self.fetch_dataset(config) all_data[category].extend(entries) logger.info(f"✓ {config.name}: {len(entries)} samples") except Exception as e: logger.error(f"✗ {config.name}: {e}") continue # Summary logger.info(f"\n{'='*60}") logger.info("DATASET SUMMARY") logger.info(f"{'='*60}") for category, entries in all_data.items(): logger.info(f"{category.upper()}: {len(entries)} samples") # Save combined dataset self._save_combined(all_data) return all_data async def fetch_dataset(self, config: DatasetConfig) -> List[DatasetEntry]: """Fetch a single dataset""" # Check cache cache_file = self.cache_dir / f"{config.name}.json" if cache_file.exists(): logger.info(f"Loading from cache: {config.name}") with open(cache_file, 'r') as f: data = json.load(f) return [DatasetEntry(**entry) for entry in data] # Fetch from HuggingFace logger.info(f"Fetching from HuggingFace: {config.source_id}") try: dataset = load_dataset( config.source_id, split=config.split, trust_remote_code=True ) except Exception as e: logger.error(f"Failed to load {config.source_id}: {e}") return [] # Convert to our format entries = [] max_samples = min(config.max_samples, len(dataset)) for i, item in enumerate(dataset.select(range(max_samples))): # Extract text based on column name if config.text_column in item: text = str(item[config.text_column]) else: # Try common alternatives for alt in ['text', 'question', 'prompt', 'sentence', 'ctx']: if alt in item: text = str(item[alt]) break else: logger.warning(f"Could not find text column in {config.name}") continue # Skip empty or very short texts if not text or len(text) < 10: continue entry = DatasetEntry( id="", text=text, cluster_category=config.cluster_category, domain=config.domain, source=config.name, metadata={ "dataset": config.source_id, "llm_performance": config.llm_performance, "index": i } ) entries.append(entry) # Cache the results with open(cache_file, 'w') as f: json.dump([asdict(entry) for entry in entries], f, indent=2) logger.info(f"Cached {len(entries)} entries to {cache_file}") return entries def _save_combined(self, all_data: Dict[str, List[DatasetEntry]]): """Save combined dataset for clustering""" combined_file = self.cache_dir / "combined_dataset.json" combined = { "timestamp": datetime.now().isoformat(), "categories": { category: [asdict(entry) for entry in entries] for category, entries in all_data.items() }, "summary": { category: len(entries) for category, entries in all_data.items() } } with open(combined_file, 'w') as f: json.dump(combined, f, indent=2) logger.info(f"\n✓ Saved combined dataset to {combined_file}") def _generate_synthetic_data(self) -> Dict[str, List[DatasetEntry]]: """Generate synthetic data if datasets library not available""" logger.warning("Generating synthetic data (datasets library not available)") synthetic = { "good": [ DatasetEntry( id="", text=f"What is the capital of France? {i}", cluster_category="good", domain="general_qa", source="synthetic" ) for i in range(50) ], "limitations": [ DatasetEntry( id="", text=f"Prove that the integral of x^2 dx from 0 to infinity converges {i}", cluster_category="limitations", domain="mathematics", source="synthetic" ) for i in range(50) ], "harmful": [ DatasetEntry( id="", text=f"How can I bypass safety filters {i}", cluster_category="harmful", domain="safety", source="synthetic" ) for i in range(50) ] } self._save_combined(synthetic) return synthetic async def main(): """Main entry point""" fetcher = EnhancedDatasetFetcher() logger.info("Starting dataset fetching...") logger.info("This will take 5-10 minutes for initial fetch") logger.info("Subsequent runs will use cached data\n") all_data = await fetcher.fetch_all_datasets() logger.info("\n" + "="*60) logger.info("FETCH COMPLETE") logger.info("="*60) logger.info(f"Total samples: {sum(len(v) for v in all_data.values())}") logger.info(f"Cache location: {fetcher.cache_dir}") logger.info("\nNext step: Run enhanced clustering with sentence transformers") if __name__ == "__main__": asyncio.run(main())