fix github issue

Author: joyful-ii-V-I

src/training/wiki_filtered_classifier_tuning/multitask_alert_filtered_wiki_classifier_training.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import os
-from collections import defaultdic
+from collections import defaultdict
 from pathlib import Path
 import random
 import argparse
@@ -14,9 +14,9 @@
 import requests
 import torch
 import torch.nn as nn
-from datasets import load_datase
-from sklearn.model_selection import train_test_spli
-from torch.utils.data import DataLoader, Datase
+from datasets import load_dataset
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader, Dataset
 from transformers import AutoModel, AutoTokenizer, get_linear_schedule_with_warmup
 import wikipediaapi
 
@@ -37,7 +37,7 @@
 
 
 WIKI_DATA_DIR = "wikipedia_data"
-ARTICLES_PER_CATEGORY = 10000
+ARTICLES_PER_CATEGORY = 10000 
 PAGEVIEWS_API_USER_AGENT = "MyMultitaskProject/1.0 ([email protected]) requests/2.0"
 
 
@@ -90,7 +90,7 @@ def __init__(self, model, tokenizer, task_configs, device="cuda"):
         self.device = device
         self.jailbreak_label_mapping = None
         self.loss_fns = {task: nn.CrossEntropyLoss() for task in task_configs if task_configs}
-
+        
         self.EXCLUDED_CATEGORIES = {
             "fiction", "culture", "arts", "comics", "media", "entertainment",
             "people", "religion", "sports", "society", "geography", "history"
@@ -104,16 +104,16 @@ def _get_wiki_articles_recursive(self, category_page, max_articles, visited_cate
         if visited_categories is None:
             visited_categories = set()
 
-        if depth >= max_depth or len(visited_categories) > max_articles * 2:
+        if depth >= max_depth or len(visited_categories) > max_articles * 2: 
             return []
-
+            
         articles = []
         visited_categories.add(category_page.title)
-
+        
         members = list(category_page.categorymembers.values())
         for member in members:
             if len(articles) >= max_articles: break
-
+            
             if member.ns == wikipediaapi.Namespace.CATEGORY and member.title not in visited_categories:
                 member_title_lower = member.title.lower()
                 if not any(excluded in member_title_lower for excluded in self.EXCLUDED_CATEGORIES):
@@ -126,7 +126,7 @@ def _get_wiki_articles_recursive(self, category_page, max_articles, visited_cate
                     ))
             elif member.ns == wikipediaapi.Namespace.MAIN:
                 articles.append(member.title)
-
+                
         return list(set(articles))
 
     def _get_pageviews(self, article_title, session, start_date, end_date):
@@ -153,7 +153,7 @@ def download_wikipedia_articles(self, categories, articles_per_category, data_di
         logger.info("--- Starting Wikipedia Data Download and Preparation ---")
         os.makedirs(data_dir, exist_ok=True)
         wiki_wiki = wikipediaapi.Wikipedia('MyMultitaskProject ([email protected])', 'en')
-
+        
         end_date = datetime.utcnow()
         start_date = end_date - timedelta(days=90)
         end_date_str = end_date.strftime('%Y%m%d')
@@ -172,9 +172,9 @@ def download_wikipedia_articles(self, categories, articles_per_category, data_di
 
             num_candidates = articles_per_category * candidate_multiplier
             logger.info(f"  Recursively searching for up to {num_candidates} candidate articles (max depth=4)...")
-
+            
             candidate_titles = self._get_wiki_articles_recursive(cat_page, num_candidates)
-
+            
             logger.info(f"  Found {len(candidate_titles)} unique candidate articles.")
 
             if not candidate_titles:
@@ -202,7 +202,7 @@ def download_wikipedia_articles(self, categories, articles_per_category, data_di
             for title in tqdm(article_titles, desc=f"  Downloading '{category}'"):
                 safe_filename = "".join([c for c in title if c.isalpha() or c.isdigit() or c==' ']).rstrip() + ".txt"
                 file_path = os.path.join(category_path, safe_filename)
-
+                
                 if os.path.exists(file_path): continue
 
                 page = wiki_wiki.page(title)
@@ -213,7 +213,7 @@ def download_wikipedia_articles(self, categories, articles_per_category, data_di
                         saved_count += 1
                     except Exception as e:
                         logger.error(f"    Could not save article '{title}': {e}")
-
+            
             logger.info(f"  Saved {saved_count} new articles for '{category}'.")
 
         logger.info("\n--- Wikipedia Data Download Complete ---")
@@ -303,7 +303,7 @@ def _load_jailbreak_dataset(self):
             # Using the 'en' (English) configuration
             jb_dataset2 = load_dataset("Babelscape/ALERT", "alert")
             texts2 = jb_dataset2['test']['prompt']
-            # Prefix labels to avoid conflicts with the other datase
+            # Prefix labels to avoid conflicts with the other dataset
             labels2 = ["ALERT:" + str(l) for l in jb_dataset2['test']['category']]
             all_texts.extend(texts2)
             all_labels.extend(labels2)
@@ -313,21 +313,21 @@ def _load_jailbreak_dataset(self):
             logger.info(f"Total loaded jailbreak samples: {len(all_texts)}")
             unique_labels = sorted(list(set(all_labels)))
             label_to_idx = {label: idx for idx, label in enumerate(unique_labels)}
-
+            
             logger.info(f"Jailbreak label distribution: {dict(sorted([(label, all_labels.count(label)) for label in set(all_labels)], key=lambda x: x[1], reverse=True))}")
             logger.info(f"Total unique jailbreak labels: {len(unique_labels)}")
-
+            
             label_indices = [label_to_idx[label] for label in all_labels]
             samples = list(zip(all_texts, label_indices))
-
+            
             self.jailbreak_label_mapping = {"label_to_idx": label_to_idx, "idx_to_label": {idx: label for label, idx in label_to_idx.items()}}
             return samples
         except Exception as e:
             import traceback
             logger.error(f"Failed to load jailbreak datasets: {e}")
             traceback.print_exc()
             return []
-
+    
     def _save_checkpoint(self, epoch, global_step, optimizer, scheduler, checkpoint_dir, label_mappings):
         os.makedirs(checkpoint_dir, exist_ok=True)
         latest_checkpoint_path = os.path.join(checkpoint_dir, 'latest_checkpoint.pt')
@@ -343,7 +343,7 @@ def _save_checkpoint(self, epoch, global_step, optimizer, scheduler, checkpoint_
         torch.save(state, latest_checkpoint_path)
         logger.info(f"Checkpoint saved for step {global_step} at {latest_checkpoint_path}")
 
-    def train(self, train_samples, val_samples, label_mappings, num_epochs=3, batch_size=16, learning_rate=2e-5,
+    def train(self, train_samples, val_samples, label_mappings, num_epochs=3, batch_size=16, learning_rate=2e-5, 
               checkpoint_dir='checkpoints', resume=False, save_steps=500, checkpoint_to_load=None):
         train_dataset = MultitaskDataset(train_samples, self.tokenizer)
         train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
@@ -372,7 +372,7 @@ def train(self, train_samples, val_samples, label_mappings, num_epochs=3, batch_
             pbar = tqdm(enumerate(train_loader), total=len(train_loader), desc=f"Epoch {epoch+1}")
             for step, batch in pbar:
                 if steps_to_skip > 0 and step < steps_to_skip: continue
-
+                
                 optimizer.zero_grad()
                 outputs = self.model(
                     input_ids=batch["input_ids"].to(self.device),
@@ -383,13 +383,13 @@ def train(self, train_samples, val_samples, label_mappings, num_epochs=3, batch_
                     task_logits = outputs[task_name][i : i + 1]
                     task_label = batch["label"][i : i + 1].to(self.device)
                     task_weight = self.task_configs[task_name].get("weight", 1.0)
-                    batch_loss += self.loss_fns[task_name](task_logits, task_label) * task_weigh
+                    batch_loss += self.loss_fns[task_name](task_logits, task_label) * task_weight
                 batch_loss.backward()
                 torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                 optimizer.step(); scheduler.step(); global_step += 1
                 if global_step > 0 and global_step % save_steps == 0:
                     self._save_checkpoint(epoch, global_step, optimizer, scheduler, checkpoint_dir, label_mappings)
-
+            
             if val_loader: self.evaluate(val_loader)
             self._save_checkpoint(epoch + 1, global_step, optimizer, scheduler, checkpoint_dir, label_mappings)
             steps_to_skip = 0
@@ -443,13 +443,13 @@ def main():
     logger.info("--- Starting Model Training ---")
 
     task_configs, label_mappings, checkpoint_to_load = {}, {}, None
-
+    
     if args.resume:
         latest_checkpoint_path = os.path.join(args.checkpoint_dir, 'latest_checkpoint.pt')
         if os.path.exists(latest_checkpoint_path):
             logger.info(f"Resuming training from checkpoint: {latest_checkpoint_path}")
             checkpoint_to_load = torch.load(latest_checkpoint_path, map_location=device)
-
+            
             task_configs = checkpoint_to_load.get('task_configs')
             label_mappings = checkpoint_to_load.get('label_mappings')
 
@@ -458,10 +458,10 @@ def main():
                 logger.warning("Cannot safely resume. Starting a fresh training run.")
                 args.resume = False
                 checkpoint_to_load = None
-                task_configs = {}
+                task_configs = {} 
             else:
                 logger.info("Loaded model configuration from checkpoint.")
-
+                
         else:
             logger.warning(f"Resume flag is set, but no checkpoint found in '{args.checkpoint_dir}'. Starting fresh run.")
             args.resume = False
@@ -476,12 +476,12 @@ def main():
             task_configs["pii"] = {"num_classes": len(label_mappings["pii"]["label_mapping"]["label_to_idx"]), "weight": 3.0}
         if "jailbreak" in label_mappings:
             task_configs["jailbreak"] = {"num_classes": len(label_mappings["jailbreak"]["label_mapping"]["label_to_idx"]), "weight": 2.0}
-
+    
     if not task_configs:
         logger.error("No tasks configured. Exiting."); return
 
     logger.info(f"Final task configurations: {task_configs}")
-
+    
     model = MultitaskBertModel(base_model_name, task_configs)
 
     if args.resume and checkpoint_to_load:
@@ -497,9 +497,9 @@ def main():
     active_label_mappings = label_mappings if (args.resume and label_mappings) else final_label_mappings
 
     trainer = MultitaskTrainer(model, tokenizer, task_configs, device)
-
+    
     logger.info(f"Total training samples: {len(train_samples)}")
-
+    
     trainer.train(
         train_samples, val_samples, active_label_mappings,
         num_epochs=10, batch_size=16,
@@ -508,7 +508,7 @@ def main():
         save_steps=args.save_steps,
         checkpoint_to_load=checkpoint_to_load
     )
-
+    
     trainer.save_model(output_path)
     with open(os.path.join(output_path, "label_mappings.json"), "w") as f:
         json.dump(active_label_mappings, f, indent=2)