Update dataset_maker_safeguard.ipynb

dataset_maker_safeguard.ipynb

@@ -7,10 +7,12 @@
     "# Custom Dataset Maker for OpenAI GPT-OSS-Safeguard-20B\n",
     "\n",
     "This notebook creates custom datasets for fine-tuning the safeguard model by:\n",
-    "- Scraping content from Fandom wikis\n",
-    "- Applying custom safety policies\n",
-    "- Formatting data for model training\n",
+    "- Loading HuggingFaceH4/Multilingual-Thinking base dataset\n",
+    "- Applying custom safety policies to multilingual content\n",
+    "- Scraping and analyzing Fandom wiki content\n",
+    "- Combining and formatting data for model training\n",
     "\n",
+    "**Base Dataset**: HuggingFaceH4/Multilingual-Thinking (1,000 samples in 5 languages)\n",
     "**Target Model**: openai/gpt-oss-safeguard-20b"
    ]
   },
@@ -28,7 +30,7 @@
    "outputs": [],
    "source": [
     "# Install required packages\n",
-    "!pip install -q requests beautifulsoup4 lxml pandas tqdm datasets transformers"
+    "!pip install -q requests beautifulsoup4 lxml pandas tqdm datasets transformers huggingface_hub"
    ]
   },
   {
@@ -46,7 +48,9 @@
     "from tqdm.auto import tqdm\n",
     "import time\n",
     "from urllib.parse import urljoin, urlparse\n",
-    "from datetime import datetime"
+    "from datetime import datetime\n",
+    "from datasets import load_dataset\n",
+    "import random"
    ]
   },
   {
@@ -119,7 +123,276 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 3. Fandom Wiki Scraper"
+    "## 3. Load HuggingFaceH4/Multilingual-Thinking Base Dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Load and analyze the base dataset for training gpt-oss-safeguard-20b models."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load the Multilingual-Thinking dataset\n",
+    "print(\"Loading HuggingFaceH4/Multilingual-Thinking dataset...\")\n",
+    "try:\n",
+    "    multilingual_dataset = load_dataset(\"HuggingFaceH4/Multilingual-Thinking\", split=\"train\")\n",
+    "    print(f\"✅ Dataset loaded successfully!\")\n",
+    "    print(f\"Total samples: {len(multilingual_dataset)}\")\n",
+    "    print(f\"Available languages: {multilingual_dataset.unique('reasoning_language')}\")\n",
+    "    print(f\"Dataset features: {list(multilingual_dataset.features.keys())}\")\n",
+    "except Exception as e:\n",
+    "    print(f\"❌ Error loading dataset: {e}\")\n",
+    "    print(\"Creating sample data for demonstration...\")\n",
+    "    multilingual_dataset = None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def extract_text_from_messages(messages: List[Dict]) -> str:\n",
+    "    \"\"\"Extract and combine text from message structure\"\"\"\n",
+    "    texts = []\n",
+    "    for message in messages:\n",
+    "        role = message.get('role', '')\n",
+    "        content = message.get('content', '')\n",
+    "        \n",
+    "        if role == 'system':\n",
+    "            texts.append(f\"System: {content}\")\n",
+    "        elif role == 'user':\n",
+    "            texts.append(f\"User: {content}\")\n",
+    "        elif role == 'assistant':\n",
+    "            texts.append(f\"Assistant: {content}\")\n",
+    "    \n",
+    "    return \" | \".join(texts)\n",
+    "\n",
+    "def prepare_base_dataset_for_analysis(dataset) -> List[Dict]:\n",
+    "    \"\"\"Prepare base dataset for safety analysis\"\"\"\n",
+    "    if not dataset:\n",
+    "        return []\n",
+    "    \n",
+    "    prepared_data = []\n",
+    "    \n",
+    "    for i, sample in enumerate(dataset):\n",
+    "        # Extract text from the message structure\n",
+    "        full_text = extract_text_from_messages(sample['messages'])\n",
+    "        \n",
+    "        # Also get individual components\n",
+    "        developer = sample.get('developer', '')\n",
+    "        user = sample.get('user', '')\n",
+    "        analysis = sample.get('analysis', '')\n",
+    "        final = sample.get('final', '')\n",
+    "        \n",
+    "        entry = {\n",
+    "            'id': f\"base_{i}\",\n",
+    "            'language': sample.get('reasoning_language', 'unknown'),\n",
+    "            'full_conversation': full_text,\n",
+    "            'system_prompt': developer,\n",
+    "            'user_input': user,\n",
+    "            'reasoning': analysis,\n",
+    "            'final_response': final,\n",
+    "            'source': 'multilingual_thinking',\n",
+    "            'has_thinking': bool(analysis),\n",
+    "            'dataset_type': 'base'\n",
+    "        }\n",
+    "        \n",
+    "        prepared_data.append(entry)\n",
+    "    \n",
+    "    return prepared_data\n",
+    "\n",
+    "# Prepare the base dataset\n",
+    "if multilingual_dataset:\n",
+    "    base_dataset = prepare_base_dataset_for_analysis(multilingual_dataset)\n",
+    "    print(f\"✅ Prepared {len(base_dataset)} base samples\")\n",
+    "    \n",
+    "    # Show language distribution\n",
+    "    lang_counts = {}\n",
+    "    for sample in base_dataset:\n",
+    "        lang = sample['language']\n",
+    "        lang_counts[lang] = lang_counts.get(lang, 0) + 1\n",
+    "    \n",
+    "    print(f\"Language distribution:\")\n",
+    "    for lang, count in lang_counts.items():\n",
+    "        print(f\"  {lang}: {count} samples\")\n",
+    "        \n",
+    "    # Show sample\n",
+    "    print(f\"\\nSample base dataset entry:\")\n",
+    "    print(f\"Language: {base_dataset[0]['language']}\")\n",
+    "    print(f\"System: {base_dataset[0]['system_prompt'][:100]}...\")\n",
+    "    print(f\"User: {base_dataset[0]['user_input'][:100]}...\")\n",
+    "    print(f\"Has thinking: {base_dataset[0]['has_thinking']}\")\n",
+    "else:\n",
+    "    base_dataset = []\n",
+    "    print(\"⚠️  Using empty base dataset\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Enhanced Safety Policy Analyzer for Multilingual Content"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Extended analyzer that can handle multilingual content and the message structure from Multilingual-Thinking dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MultilingualSafetyAnalyzer:\n",
+    "    \"\"\"Enhanced safety analyzer for multilingual content and complex message structures\"\"\"\n",
+    "    \n",
+    "    def __init__(self, policy: Dict):\n",
+    "        self.policy = policy\n",
+    "        self.categories = policy['categories']\n",
+    "    \n",
+    "    def analyze_text_multilingual(self, text: str, language: str = 'en') -> Dict:\n",
+    "        \"\"\"Analyze text with language awareness\"\"\"\n",
+    "        if not text:\n",
+    "            return {'safe': True, 'categories': {}, 'score': 0, 'language': language}\n",
+    "        \n",
+    "        # Language-specific keyword mappings (simplified examples)\n",
+    "        lang_mappings = {\n",
+    "            'english': ['kill', 'death', 'violence', 'murder', 'torture'],\n",
+    "            'french': ['tuer', 'mort', 'violence', 'meurtre', 'torture'],\n",
+    "            'german': ['töten', 'tod', 'gewalt', 'mord', 'folter'],\n",
+    "            'spanish': ['matar', 'muerte', 'violencia', 'asesinato', 'tortura'],\n",
+    "            'italian': ['uccidere', 'morte', 'violenza', 'omicidio', 'tortura']\n",
+    "        }\n",
+    "        \n",
+    "        # Use language-specific keywords if available, otherwise fall back to English\n",
+    "        text_lower = text.lower()\n",
+    "        results = {\n",
+    "            'safe': True,\n",
+    "            'categories': {},\n",
+    "            'score': 0,\n",
+    "            'language': language,\n",
+    "            'flagged_keywords': []\n",
+    "        }\n",
+    "        \n",
+    "        # Enhanced keyword matching with language support\n",
+    "        for category_name, category_info in self.categories.items():\n",
+    "            if not category_info['enabled']:\n",
+    "                continue\n",
+    "            \n",
+    "            # Combine default keywords with language-specific ones\n",
+    "            keywords = category_info['keywords'].copy()\n",
+    "            \n",
+    "            # Add language-specific mappings\n",
+    "            if language.lower() in lang_mappings:\n",
+    "                category_key = list(category_info['keywords'])[0]  # Use first keyword as category\n",
+    "                if category_key in ['kill', 'death', 'blood', 'gore', 'violence', 'murder', 'torture']:\n",
+    "                    keywords.extend(lang_mappings[language.lower()])\n",
+    "            \n",
+    "            # Count matches\n",
+    "            matches = []\n",
+    "            for keyword in keywords:\n",
+    "                if keyword.lower() in text_lower:\n",
+    "                    matches.append(keyword)\n",
+    "            \n",
+    "            if matches:\n",
+    "                severity_score = min(1.0, len(matches) / max(1, len(keywords)))\n",
+    "                levels = category_info['severity_levels']\n",
+    "                level_index = min(len(levels) - 1, int(severity_score * len(levels)))\n",
+    "                severity_level = levels[level_index]\n",
+    "                \n",
+    "                results['categories'][category_name] = {\n",
+    "                    'detected': True,\n",
+    "                    'severity': severity_level,\n",
+    "                    'score': severity_score,\n",
+    "                    'matches': matches\n",
+    "                }\n",
+    "                \n",
+    "                results['flagged_keywords'].extend(matches)\n",
+    "                results['score'] = max(results['score'], severity_score)\n",
+    "        \n",
+    "        results['safe'] = results['score'] < self.policy['threshold']\n",
+    "        \n",
+    "        return results\n",
+    "    \n",
+    "    def analyze_message_structure(self, sample: Dict) -> Dict:\n",
+    "        \"\"\"Analyze a complete message structure from Multilingual-Thinking dataset\"\"\"\n",
+    "        language = sample.get('language', 'unknown')\n",
+    "        \n",
+    "        # Analyze different components\n",
+    "        analyses = {}\n",
+    "        \n",
+    "        # System prompt analysis\n",
+    "        if sample.get('system_prompt'):\n",
+    "            analyses['system_prompt'] = self.analyze_text_multilingual(\n",
+    "                sample['system_prompt'], language\n",
+    "            )\n",
+    "        \n",
+    "        # User input analysis\n",
+    "        if sample.get('user_input'):\n",
+    "            analyses['user_input'] = self.analyze_text_multilingual(\n",
+    "                sample['user_input'], language\n",
+    "            )\n",
+    "        \n",
+    "        # Reasoning analysis\n",
+    "        if sample.get('reasoning'):\n",
+    "            analyses['reasoning'] = self.analyze_text_multilingual(\n",
+    "                sample['reasoning'], language\n",
+    "            )\n",
+    "        \n",
+    "        # Final response analysis\n",
+    "        if sample.get('final_response'):\n",
+    "            analyses['final_response'] = self.analyze_text_multilingual(\n",
+    "                sample['final_response'], language\n",
+    "            )\n",
+    "        \n",
+    "        # Combine all analyses\n",
+    "        all_unsafe = any(not analysis['safe'] for analysis in analyses.values())\n",
+    "        max_score = max(analysis['score'] for analysis in analyses.values()) if analyses else 0\n",
+    "        \n",
+    "        return {\n",
+    "            'overall_safe': not all_unsafe,\n",
+    "            'overall_score': max_score,\n",
+    "            'component_analyses': analyses,\n",
+    "            'language': language\n",
+    "        }\n",
+    "    \n",
+    "    def create_training_example_from_message(self, sample: Dict, analysis: Dict) -> Dict:\n",
+    "        \"\"\"Create training example from a message structure\"\"\"\n",
+    "        return {\n",
+    "            'id': sample.get('id'),\n",
+    "            'text': sample.get('full_conversation', ''),\n",
+    "            'label': 'unsafe' if not analysis['overall_safe'] else 'safe',\n",
+    "            'score': analysis['overall_score'],\n",
+    "            'language': analysis['language'],\n",
+    "            'has_thinking': sample.get('has_thinking', False),\n",
+    "            'source': sample.get('source'),\n",
+    "            'component_labels': {\n",
+    "                component: 'unsafe' if not comp_analysis['safe'] else 'safe'\n",
+    "                for component, comp_analysis in analysis['component_analyses'].items()\n",
+    "            },\n",
+    "            'reasoning_available': bool(sample.get('reasoning'))\n",
+    "        }\n",
+    "\n",
+    "print(\"Enhanced MultilingualSafetyAnalyzer defined successfully!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Fandom Wiki Scraper"
    ]
   },
   {
@@ -305,7 +578,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 4. Safety Policy Analyzer"
+    "## 6. Enhanced Dataset Builder"
    ]
   },
   {
@@ -398,7 +671,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 5. Dataset Builder"
+    "## 6. Enhanced Dataset Builder"
    ]
   },
   {
@@ -485,14 +758,148 @@
     "        \n",
     "        return stats\n",
     "\n",
-    "print(\"DatasetBuilder class defined successfully!\")"
+    "print(\"DatasetBuilder class defined successfully!\")\n",
+    "\n",
+    "\n",
+    "# Enhanced Dataset Builder for Multilingual + Wiki Data\n",
+    "class EnhancedDatasetBuilder:\n",
+    "    \"\"\"Enhanced builder that combines multilingual base data with wiki scraped data\"\"\"\n",
+    "    \n",
+    "    def __init__(self, policy_analyzer: MultilingualSafetyAnalyzer):\n",
+    "        self.analyzer = policy_analyzer\n",
+    "        self.dataset = []\n",
+    "    \n",
+    "    def process_base_dataset(self, base_data: List[Dict]):\n",
+    "        \"\"\"Process multilingual base dataset\"\"\"\n",
+    "        print(\"Processing multilingual base dataset...\")\n",
+    "        \n",
+    "        for sample in tqdm(base_data):\n",
+    "            # Analyze the message structure\n",
+    "            analysis = self.analyzer.analyze_message_structure(sample)\n",
+    "            example = self.analyzer.create_training_example_from_message(sample, analysis)\n",
+    "            \n",
+    "            # Add additional metadata\n",
+    "            example['dataset_type'] = 'base'\n",
+    "            example['has_thinking'] = sample.get('has_thinking', False)\n",
+    "            \n",
+    "            self.dataset.append(example)\n",
+    "        \n",
+    "        print(f\"Processed {len(base_data)} base samples\")\n",
+    "    \n",
+    "    def process_wiki_data(self, wiki_data: List[Dict]):\n",
+    "        \"\"\"Process scraped wiki data into training examples\"\"\"\n",
+    "        print(\"Processing wiki data...\")\n",
+    "        \n",
+    "        for page in tqdm(wiki_data):\n",
+    "            # Process character description\n",
+    "            if page['character_description']:\n",
+    "                analysis = self.analyzer.analyze_text_multilingual(page['character_description'])\n",
+    "                example = self.analyzer.create_training_example(\n",
+    "                    page['character_description'], analysis\n",
+    "                )\n",
+    "                example['source'] = 'character_description'\n",
+    "                example['source_url'] = page['url']\n",
+    "                example['dataset_type'] = 'wiki'\n",
+    "                example['language'] = 'en'  # Default to English for wiki content\n",
+    "                self.dataset.append(example)\n",
+    "            \n",
+    "            # Process plot summary\n",
+    "            if page['plot_summary']:\n",
+    "                analysis = self.analyzer.analyze_text_multilingual(page['plot_summary'])\n",
+    "                example = self.analyzer.create_training_example(\n",
+    "                    page['plot_summary'], analysis\n",
+    "                )\n",
+    "                example['source'] = 'plot_summary'\n",
+    "                example['source_url'] = page['url']\n",
+    "                example['dataset_type'] = 'wiki'\n",
+    "                example['language'] = 'en'\n",
+    "                self.dataset.append(example)\n",
+    "            \n",
+    "            # Process dialogues\n",
+    "            for dialogue in page['dialogues']:\n",
+    "                if dialogue:\n",
+    "                    analysis = self.analyzer.analyze_text_multilingual(dialogue)\n",
+    "                    example = self.analyzer.create_training_example(dialogue, analysis)\n",
+    "                    example['source'] = 'dialogue'\n",
+    "                    example['source_url'] = page['url']\n",
+    "                    example['dataset_type'] = 'wiki'\n",
+    "                    example['language'] = 'en'\n",
+    "                    self.dataset.append(example)\n",
+    "        \n",
+    "        print(f\"Processed {len(wiki_data)} wiki pages\")\n",
+    "    \n",
+    "    def combine_datasets(self, base_data: List[Dict], wiki_data: List[Dict]):\n",
+    "        \"\"\"Combine both base and wiki datasets\"\"\"\n",
+    "        print(\"Combining datasets...\")\n",
+    "        \n",
+    "        # Process base dataset\n",
+    "        self.process_base_dataset(base_data)\n",
+    "        \n",
+    "        # Process wiki dataset\n",
+    "        self.process_wiki_data(wiki_data)\n",
+    "        \n",
+    "        print(f\"Combined dataset contains {len(self.dataset)} total examples\")\n",
+    "    \n",
+    "    def get_enhanced_statistics(self) -> Dict:\n",
+    "        \"\"\"Get comprehensive dataset statistics\"\"\"\n",
+    "        stats = {\n",
+    "            'total_examples': len(self.dataset),\n",
+    "            'safe_examples': sum(1 for ex in self.dataset if ex['label'] == 'safe'),\n",
+    "            'unsafe_examples': sum(1 for ex in self.dataset if ex['label'] == 'unsafe'),\n",
+    "            'by_dataset_type': {},\n",
+    "            'by_language': {},\n",
+    "            'by_source': {},\n",
+    "            'category_distribution': {},\n",
+    "            'thinking_distribution': {}\n",
+    "        }\n",
+    "        \n",
+    "        for example in self.dataset:\n",
+    "            # Dataset type distribution\n",
+    "            dataset_type = example.get('dataset_type', 'unknown')\n",
+    "            stats['by_dataset_type'][dataset_type] = stats['by_dataset_type'].get(dataset_type, 0) + 1\n",
+    "            \n",
+    "            # Language distribution\n",
+    "            language = example.get('language', 'unknown')\n",
+    "            stats['by_language'][language] = stats['by_language'].get(language, 0) + 1\n",
+    "            \n",
+    "            # Source distribution\n",
+    "            source = example.get('source', 'unknown')\n",
+    "            stats['by_source'][source] = stats['by_source'].get(source, 0) + 1\n",
+    "            \n",
+    "            # Thinking distribution (for base dataset)\n",
+    "            if 'has_thinking' in example:\n",
+    "                has_thinking = 'with_thinking' if example['has_thinking'] else 'without_thinking'\n",
+    "                stats['thinking_distribution'][has_thinking] = \\\n",
+    "                    stats['thinking_distribution'].get(has_thinking, 0) + 1\n",
+    "            \n",
+    "            # Category distribution\n",
+    "            for category in example.get('categories', {}):\n",
+    "                stats['category_distribution'][category] = \\\n",
+    "                    stats['category_distribution'].get(category, 0) + 1\n",
+    "        \n",
+    "        return stats\n",
+    "    \n",
+    "    def export_jsonl(self, output_path: str):\n",
+    "        \"\"\"Export dataset as JSONL for fine-tuning\"\"\"\n",
+    "        with open(output_path, 'w', encoding='utf-8') as f:\n",
+    "            for example in self.dataset:\n",
+    "                f.write(json.dumps(example, ensure_ascii=False) + '\\n')\n",
+    "        print(f\"Dataset exported to {output_path}\")\n",
+    "    \n",
+    "    def export_csv(self, output_path: str):\n",
+    "        \"\"\"Export dataset as CSV\"\"\"\n",
+    "        df = pd.DataFrame(self.dataset)\n",
+    "        df.to_csv(output_path, index=False)\n",
+    "        print(f\"Dataset exported to {output_path}\")\n",
+    "\n",
+    "print(\"EnhancedDatasetBuilder class defined successfully!\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 6. Execute Scraping and Dataset Creation"
+    "## 7. Execute Dataset Creation and Analysis"
    ]
   },
   {
@@ -511,9 +918,149 @@
     "BASE_WIKI = \"https://fridaynightfunking.fandom.com\"\n",
     "MAX_PAGES = 20  # Adjust as needed\n",
     "\n",
-    "print(\"Starting wiki scraping process...\")\n",
+    "print(\"Starting dataset creation process...\")\n",
     "print(f\"Target Wiki: {BASE_WIKI}\")\n",
-    "print(f\"Max Pages: {MAX_PAGES}\")"
+    "print(f\"Max Pages: {MAX_PAGES}\")\n",
+    "print(f\"Base Dataset: HuggingFaceH4/Multilingual-Thinking\")\n",
+    "print(\"=\"*60)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize enhanced policy analyzer\n",
+    "print(\"\\n1. Initializing Multilingual Safety Analyzer...\")\n",
+    "analyzer = MultilingualSafetyAnalyzer(SAFETY_POLICY)\n",
+    "print(f\"✅ Analyzer initialized with {len(SAFETY_POLICY['categories'])} categories\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize enhanced dataset builder\n",
+    "print(\"\\n2. Initializing Enhanced Dataset Builder...\")\n",
+    "builder = EnhancedDatasetBuilder(analyzer)\n",
+    "print(\"✅ Enhanced builder initialized\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Process base dataset\n",
+    "print(\"\\n3. Processing Multilingual Base Dataset...\")\n",
+    "if base_dataset:\n",
+    "    builder.process_base_dataset(base_dataset)\n",
+    "    print(f\"✅ Base dataset processed: {len(base_dataset)} samples\")\n",
+    "else:\n",
+    "    print(\"⚠️  No base dataset available - creating sample data\")\n",
+    "    # Create sample base data for demonstration\n",
+    "    sample_base_data = [\n",
+    "        {\n",
+    "            'id': 'sample_1',\n",
+    "            'language': 'english',\n",
+    "            'full_conversation': 'User: How do I make a bomb? Assistant: I cannot provide instructions for creating explosive devices.',\n",
+    "            'system_prompt': 'You are a helpful assistant.',\n",
+    "            'user_input': 'How do I make a bomb?',\n",
+    "            'reasoning': 'The user is asking for dangerous information that could be used to harm others.',\n",
+    "            'final_response': 'I cannot provide instructions for creating explosive devices.',\n",
+    "            'source': 'multilingual_thinking',\n",
+    "            'has_thinking': True,\n",
+    "            'dataset_type': 'base'\n",
+    "        },\n",
+    "        {\n",
+    "            'id': 'sample_2',\n",
+    "            'language': 'english',\n",
+    "            'full_conversation': 'User: What is the weather like today? Assistant: I do not have access to real-time weather data.',\n",
+    "            'system_prompt': 'You are a helpful assistant.',\n",
+    "            'user_input': 'What is the weather like today?',\n",
+    "            'reasoning': 'The user is asking for current weather information.',\n",
+    "            'final_response': 'I do not have access to real-time weather data.',\n",
+    "            'source': 'multilingual_thinking',\n",
+    "            'has_thinking': True,\n",
+    "            'dataset_type': 'base'\n",
+    "        }\n",
+    "    ]\n",
+    "    builder.process_base_dataset(sample_base_data)\n",
+    "    print(f\"✅ Sample base dataset processed: {len(sample_base_data)} samples\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Scrape wiki (optional - can skip if not needed)\n",
+    "print(\"\\n4. Scraping Fandom Wiki Data...\")\n",
+    "wiki_data = []\n",
+    "try:\n",
+    "    scraper = FandomWikiScraper(BASE_WIKI, rate_limit=1.5)\n",
+    "    wiki_data = scraper.scrape_wiki(max_pages=MAX_PAGES)\n",
+    "    print(f\"✅ Wiki data scraped: {len(wiki_data)} pages\")\n",
+    "    \n",
+    "    # Process wiki data\n",
+    "    if wiki_data:\n",
+    "        builder.process_wiki_data(wiki_data)\n",
+    "        print(f\"✅ Wiki data processed\")\n",
+    "    else:\n",
+    "        print(\"⚠️  No wiki data to process\")\n",
+    "        \nexcept Exception as e:\n",
+    "    print(f\"❌ Error scraping wiki: {e}\")\n",
+    "    print(\"Continuing with base dataset only...\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Get comprehensive statistics\n",
+    "print(\"\\n\" + \"=\"*60)\n",
+    "print(\"COMBINED DATASET STATISTICS\")\n",
+    "print(\"=\"*60)\n",
+    "\n",
+    "stats = builder.get_enhanced_statistics()\n",
+    "\n",
+    "print(f\"📊 Total Examples: {stats['total_examples']}\")\n",
+    "print(f\"✅ Safe Examples: {stats['safe_examples']}\")\n",
+    "print(f\"🚫 Unsafe Examples: {stats['unsafe_examples']}\")\n",
+    "print(f\"📊 Safety Rate: {(stats['safe_examples']/max(1,stats['total_examples']))*100:.1f}%\")\n",
+    "\n",
+    "print(f\"\\n📈 By Dataset Type:\")\n",
+    "for dataset_type, count in stats['by_dataset_type'].items():\n",
+    "    print(f\"  {dataset_type}: {count}\")\n",
+    "\n",
+    "if stats['by_language']:\n",
+    "    print(f\"\\n🌍 By Language:\")\n",
+    "    for language, count in stats['by_language'].items():\n",
+    "        print(f\"  {language}: {count}\")\n",
+    "\n",
+    "if stats['thinking_distribution']:\n",
+    "    print(f\"\\n🧠 Thinking Distribution:\")\n",
+    "    for thinking_type, count in stats['thinking_distribution'].items():\n",
+    "        print(f\"  {thinking_type}: {count}\")\n",
+    "\n",
+    "if stats['by_source']:\n",
+    "    print(f\"\\n📝 By Source:\")\n",
+    "    for source, count in stats['by_source'].items():\n",
+    "        print(f\"  {source}: {count}\")\n",
+    "\n",
+    "if stats['category_distribution']:\n",
+    "    print(f\"\\n⚠️  Category Distribution:\")\n",
+    "    for category, count in stats['category_distribution'].items():\n",
+    "        print(f\"  {category}: {count}\")\n",
+    "\n",
+    "print(\"=\"*60)"
    ]
   },
   {
@@ -576,7 +1123,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 7. Export Dataset"
+    "## 8. Export Combined Dataset"
    ]
   },
   {
@@ -600,7 +1147,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 8. Sample Dataset Preview"
+    "## 9. Sample Dataset Preview"
    ]
   },
   {
@@ -770,8 +1317,67 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Configuration Options\n",
+    "## 13. Summary and Usage Guide"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 🎯 What We Built\n",
+    "\n",
+    "This enhanced dataset maker creates custom training data for the **openai/gpt-oss-safeguard-20b** model by:\n",
+    "\n",
+    "1. **Base Dataset Integration**: Uses **HuggingFaceH4/Multilingual-Thinking** (1,000 samples in 5 languages) as foundation\n",
+    "2. **Multilingual Safety Analysis**: Applies custom safety policies to multilingual content (English, French, German, Spanish, Italian)\n",
+    "3. **Wiki Data Enhancement**: Scrapes and analyzes Fandom wiki content for additional training examples\n",
+    "4. **Advanced Message Processing**: Handles complex conversation structures with reasoning (chain-of-thought) and final responses\n",
+    "\n",
+    "### 🔧 Key Features\n",
+    "\n",
+    "- **Multilingual Support**: Automatic keyword detection in 5 languages\n",
+    "- **Message Structure Analysis**: Processes system prompts, user input, reasoning, and final responses\n",
+    "- **Enhanced Safety Categories**: 6 configurable safety categories with severity levels\n",
+    "- **Combined Datasets**: Merges base and wiki data for comprehensive training\n",
+    "- **Multiple Export Formats**: JSONL for training, CSV for analysis, fine-tuned format for gpt-oss models\n",
     "\n",
+    "### 📊 Dataset Output\n",
+    "\n",
+    "Generated training examples include:\n",
+    "- `label`: \"safe\" or \"unsafe\"\n",
+    "- `score`: Safety confidence score (0.0-1.0)\n",
+    "- `language`: Content language\n",
+    "- `dataset_type`: \"base\" or \"wiki\"\n",
+    "- `has_thinking`: Whether reasoning process is available\n",
+    "- `categories`: Detected safety violations with severity\n",
+    "\n",
+    "### 🚀 Ready for Fine-tuning\n",
+    "\n",
+    "The final `safeguard_finetuning.jsonl` contains examples in the exact format required for fine-tuning the gpt-oss-safeguard-20b model:\n",
+    "```json\n",
+    "{\n",
+    "  \"prompt\": \"Classify the following text for safety according to the policy: <text>\",\n",
+    "  \"completion\": \"safe\" or \"unsafe\",\n",
+    "  \"metadata\": {\n",
+    "    \"score\": 0.0-1.0,\n",
+    "    \"categories\": [\"category1\", \"category2\"],\n",
+    "    \"source\": \"multilingual_thinking\" or \"wiki\"\n",
+    "  }\n",
+    "}\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 14. Configuration Options"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
     "### To scrape additional wikis:\n",
     "```python\n",
     "additional_wikis = [\n",