Create main.py

HROM-V1.6/main.py

@@ -0,0 +1,1428 @@
+import os
+# Set parallelism env var *before* importing tokenizers
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+# Import necessary dataset functions, including concatenate_datasets if needed later
+from datasets import load_dataset, disable_caching, concatenate_datasets
+from tokenizers import Tokenizer, models, trainers, pre_tokenizers, processors, decoders
+import math
+import re
+from datetime import datetime
+from contextlib import nullcontext
+from collections import defaultdict
+import logging
+import random # For shuffling combined data
+
+# Disable caching for datasets if needed, helps ensure reprocessing
+# disable_caching()
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    force=True  # Add this
+)
+
+# Configuration
+CONFIG = {
+    # --- Scaled Parameters ---
+    "dim": 960,
+    "n_layers": 32,
+    "n_heads": 32,
+    "ff_dim": 3608, # Explicitly set to 4 * dim
+
+    # --- Kept Parameters ---
+    "dropout": 0.1,
+    "max_seq_len": 512,
+    "vocab_size": 32000, # Fixed by tokenizer
+
+    # --- Training/Dataset Parameters ---
+    "batch_size": 2,  # Increased batch size for 2 GPUs (adjust if needed)
+    "checkpoint_interval": 2000,
+    "debug_interval": 400,
+    # --- ADDED CoQA and QuAC and Universal-Transformers-RawConvo ---
+    "datasets": ["daily_dialog", "empathetic_dialogues", "blended_skill_talk", "AlekseyKorshuk/persona-chat", "elapt1c/Universal-Transformers-RawConvo", "papahawk/conversational-01"],
+    "tokenizer_name": "hrom_tokenizer.json", # New name for expanded tokenizer
+    "checkpoint_dir": "checkpoints", # Separate directory for expanded data model
+    # --- Increased samples per dataset slightly for tokenizer ---
+    "tokenizer_train_samples_per_dataset": 100000, # Use same limit for all, incl. new ones
+    "learning_rate": 1e-5,
+    "warmup_steps": 1000,
+    "max_turns": 8, # Keep max_turns limit for Q&A datasets too
+    "max_checkpoints": 5,
+    "num_epochs": 30,
+    "grad_accum_steps": 4,
+    "num_gpus": 2 # Specify the number of GPUs to use
+}
+
+# --- Model Definition (HROM, HROMBlock, HROMAttention, SwiGLU, RoPE) ---
+# (These classes remain unchanged from the previous version)
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+    def forward(self, seq_len):
+        t = torch.arange(seq_len, device=self.inv_freq.device).type_as(self.inv_freq)
+        freqs = torch.einsum("i, j -> i j", t, self.inv_freq)
+        if seq_len == 0:
+             return torch.empty((0, self.inv_freq.shape[0] * 2), device=self.inv_freq.device)
+        # Defensive reshape only if necessary
+        if freqs.shape[0] != seq_len and seq_len > 0:
+             freqs = freqs.reshape(seq_len, -1)
+        elif seq_len == 0: # Handle edge case for empty sequences
+            return torch.empty((0, self.inv_freq.shape[0]*2), device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+
+        return torch.cat((freqs, freqs), dim=-1)
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(pos, t):
+    # pos: (T, dim_rotary), t: (B, H, T, Head_Dim)
+    pos = pos.to(t.device, dtype=t.dtype)
+    pos = pos.unsqueeze(0).unsqueeze(1) # Shape: (1, 1, T, dim_rotary)
+    tensor_seq_len = t.shape[2]
+    pos_seq_len = pos.shape[2]
+
+    if pos_seq_len < tensor_seq_len:
+         logging.warning(f"RoPE Warning: pos sequence length ({pos_seq_len}) is shorter than tensor sequence length ({tensor_seq_len}). Using truncated tensor length for RoPE.")
+         # This case is tricky, maybe only apply to the length of pos?
+         # Or indicates an issue upstream. Let's slice t for now, though it's unusual.
+         t_rotated = t[:, :, :pos_seq_len, :]
+         pos = pos[:, :, :pos_seq_len, :] # Ensure pos matches the sliced tensor length
+
+         # Apply rotation only to the slice
+         cos_pos = pos.cos()
+         sin_pos = pos.sin()
+         t_rotated = (t_rotated * cos_pos) + (rotate_half(t_rotated) * sin_pos)
+
+         # Concatenate the rotated part with the un-rotated part
+         t_unrotated = t[:, :, pos_seq_len:, :]
+         return torch.cat([t_rotated, t_unrotated], dim=2)
+
+    elif pos_seq_len > tensor_seq_len:
+         pos = pos[:, :, :tensor_seq_len, :] # Slice pos to match tensor
+
+    # Check dimension match after potential slicing
+    if pos.shape[-1] != t.shape[-1]:
+        logging.error(f"Mismatched dimensions for RoPE: pos ({pos.shape[-1]}) vs t ({t.shape[-1]})")
+        raise ValueError("Rotary embedding dimension must match head dimension.")
+
+    cos_pos = pos.cos()
+    sin_pos = pos.sin()
+    rotated_t = (t * cos_pos) + (rotate_half(t) * sin_pos)
+    return rotated_t
+
+
+class SwiGLU(nn.Module):
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1)
+        return x * nn.functional.gelu(gate)
+
+class HROMAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dim = CONFIG["dim"]
+        self.n_heads = CONFIG["n_heads"]
+        self.head_dim = self.dim // self.n_heads
+        if self.dim % self.n_heads != 0:
+            raise ValueError("dim must be divisible by n_heads")
+        self.qkv = nn.Linear(self.dim, 3 * self.dim)
+        self.proj = nn.Linear(self.dim, self.dim)
+        # Initialize RotaryEmbedding with head_dim to ensure match
+        self.rotary = RotaryEmbedding(self.head_dim)
+        self.dropout = nn.Dropout(CONFIG["dropout"])
+
+        # --- Assertion to verify dimension match ---
+        assert self.rotary.inv_freq.shape[0] * 2 == self.head_dim, f"Rotary Embedding dimension ({self.rotary.inv_freq.shape[0] * 2}) must match head dimension ({self.head_dim}). Check CONFIG parameters 'dim' and 'n_heads'."
+
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape
+        qkv = self.qkv(x)
+        qkv = qkv.reshape(B, T, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(2)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        # Generate RoPE embeddings for the current sequence length T
+        pos = self.rotary(T) # Shape (T, Head_Dim)
+        # Apply RoPE
+        q = apply_rotary_pos_emb(pos, q)
+        k = apply_rotary_pos_emb(pos, k)
+        # Attention calculation
+        attn_scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
+        if mask is not None:
+            # Ensure mask is broadcastable (B, 1, T, T)
+            if mask.dim() == 2: # (B, T) -> (B, 1, 1, T) -> add with causal = (B, 1, T, T)
+                mask = mask.unsqueeze(1).unsqueeze(2)
+            elif mask.dim() == 3: # (B, T, T)
+                mask = mask.unsqueeze(1)
+            # Add mask AFTER scaling scores
+            attn_scores = attn_scores + mask # Add large negative values for masked positions
+        # Softmax and dropout
+        attn_probs = torch.softmax(attn_scores.float(), dim=-1).to(dtype=x.dtype) # Use float for stability
+        attn_probs = self.dropout(attn_probs)
+        # Output projection
+        output = attn_probs @ v
+        output = output.transpose(1, 2).reshape(B, T, self.dim)
+        return self.proj(output)
+
+
+class HROMBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = HROMAttention()
+        self.ff = nn.Sequential(
+            nn.Linear(CONFIG["dim"], 2 * CONFIG["ff_dim"]),
+            SwiGLU(),
+            nn.Linear(CONFIG["ff_dim"], CONFIG["dim"])
+        )
+        self.norm1 = nn.LayerNorm(CONFIG["dim"])
+        self.norm2 = nn.LayerNorm(CONFIG["dim"])
+        self.dropout = nn.Dropout(CONFIG["dropout"])
+
+    def forward(self, x, mask=None):
+        # Pre-Normalization
+        normed_x = self.norm1(x)
+        attn_output = self.attn(normed_x, mask)
+        x = x + self.dropout(attn_output)
+
+        normed_x = self.norm2(x)
+        ff_output = self.ff(normed_x)
+        x = x + self.dropout(ff_output)
+        return x
+
+class HROM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed = nn.Embedding(CONFIG["vocab_size"], CONFIG["dim"])
+        self.blocks = nn.ModuleList([HROMBlock() for _ in range(CONFIG["n_layers"])])
+        self.norm = nn.LayerNorm(CONFIG["dim"])
+        self.head = nn.Linear(CONFIG["dim"], CONFIG["vocab_size"])
+        self.dropout = nn.Dropout(CONFIG["dropout"]) # Add dropout after embedding
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.xavier_uniform_(module.weight)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+             torch.nn.init.zeros_(module.bias)
+             torch.nn.init.ones_(module.weight)
+
+    def forward(self, input_ids, attention_mask=None):
+        B, T = input_ids.shape
+        x = self.embed(input_ids)
+        x = self.dropout(x) # Apply dropout after embedding
+
+        # Create the combined mask for attention
+        combined_mask = None
+        # Start with causal mask valid for all sequences in batch
+        causal_mask = torch.triu(torch.ones(T, T, device=input_ids.device) * float('-inf'), diagonal=1)
+        combined_mask = causal_mask.unsqueeze(0).unsqueeze(1) # (1, 1, T, T)
+
+        if attention_mask is not None:
+            # Process padding mask from attention_mask (0 = pad, 1 = real)
+            # Convert 0s to -inf, 1s to 0
+            pad_mask = (1.0 - attention_mask.to(torch.float32)) * torch.finfo(torch.float32).min
+            pad_mask = pad_mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, T)
+            # Add padding mask to causal mask. Broadcasting ensures (B, 1, T, T)
+            # Where pad_mask is -inf, the result is -inf. Otherwise, it's the causal value.
+            combined_mask = combined_mask + pad_mask
+
+        # Ensure mask dtype matches data dtype (esp. for AMP)
+        combined_mask = combined_mask.to(dtype=x.dtype)
+
+        for block in self.blocks:
+            x = block(x, combined_mask) # Pass the combined mask to each block
+
+        x = self.norm(x)
+        logits = self.head(x)
+        return logits
+
+# --- Tokenizer Training ---
+
+class TokenizerTrainer:
+    def __init__(self):
+        self.tokenizer = Tokenizer(models.BPE())
+        self.tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+        self.tokenizer.decoder = decoders.ByteLevel()
+        self.special_tokens = ["<pad>", "<s>", "</s>", "<unk>", "<user>", "<assistant>"]
+        # Use the updated tokenizer name from CONFIG
+        self.tokenizer_path = os.path.join("tokenizer", CONFIG["tokenizer_name"])
+        self.tokenizer_dir = os.path.dirname(self.tokenizer_path)
+
+    def _clean_text(self, text):
+        text = str(text) # Ensure text is string
+        text = re.sub(r'_comma_', ',', text)
+        # Allow alphanumeric, whitespace, and basic punctuation including quotes
+        text = re.sub(r'[^\w\s.,!?\'\-:;<>"]', '', text)
+        text = re.sub(r'\s+', ' ', text).strip()
+        return text
+
+    def train(self, dataset_names):
+        logging.info("Starting tokenizer training...")
+        text_samples = []
+        samples_per_dataset = CONFIG['tokenizer_train_samples_per_dataset']
+
+        # --- Process DailyDialog ---
+        if "daily_dialog" in dataset_names:
+            logging.info(f"Loading daily_dialog for tokenizer training (max {samples_per_dataset} dialogues)...")
+            try:
+                # Limit dialogues loaded directly using slicing
+                dd_dataset = load_dataset("daily_dialog", split=f"train[:{samples_per_dataset}]", trust_remote_code=True) # Add trust_remote_code=True
+                logging.info("Processing daily_dialog...")
+                for entry in dd_dataset:
+                    formatted_dialogue = []
+                    dialogue = entry['dialog'][:CONFIG["max_turns"]]
+                    for i, utterance in enumerate(dialogue):
+                        role = "<user>" if i % 2 == 0 else "<assistant>"
+                        cleaned_utterance = self._clean_text(utterance)
+                        if cleaned_utterance: # Only add non-empty turns
+                             formatted_dialogue.append(f"{role} {cleaned_utterance}")
+                    if formatted_dialogue: # Only add if dialogue is not empty after cleaning
+                        text_samples.append(" </s> ".join(formatted_dialogue))
+            except Exception as e:
+                logging.error(f"Failed to load or process daily_dialog for tokenizer: {e}")
+
+        # --- Process EmpatheticDialogues ---
+        if "empathetic_dialogues" in dataset_names:
+            logging.info(f"Loading empathetic_dialogues for tokenizer training (max {samples_per_dataset} dialogues)...")
+            try:
+                # Load more initially to ensure we get enough unique conversations (adjust multiplier if needed)
+                ed_dataset = load_dataset("empathetic_dialogues", split=f"train[:{samples_per_dataset * 3}]", trust_remote_code=True) # Add trust_remote_code=True
+                logging.info("Processing empathetic_dialogues...")
+                conversations = defaultdict(list)
+                processed_conv_count = 0
+                # Group utterances by conv_id first
+                grouped_by_conv = defaultdict(list)
+                for entry in ed_dataset:
+                    grouped_by_conv[entry['conv_id']].append(entry)
+
+                # Process conversations ensuring max samples limit
+                for conv_id, entries in grouped_by_conv.items():
+                    if processed_conv_count >= samples_per_dataset:
+                        break
+                    # Sort by utterance_idx to maintain order
+                    sorted_entries = sorted(entries, key=lambda x: x['utterance_idx'])
+                    formatted_dialogue = []
+                    # Handle context and first utterance
+                    if sorted_entries[0]['context']:
+                         cleaned_context = self._clean_text(sorted_entries[0]['context'])
+                         if cleaned_context:
+                              formatted_dialogue.append(f"<user> {cleaned_context}") # Assume context is user start
+                    # Process subsequent utterances
+                    last_role = '<user>' if formatted_dialogue else None # Set initial last role based on context
+                    for entry in sorted_entries:
+                        cleaned_utterance = self._clean_text(entry['utterance'])
+                        if cleaned_utterance:
+                            # Determine role based on alternation
+                            current_role = '<assistant>' if last_role == '<user>' else '<user>'
+                            formatted_dialogue.append(f"{current_role} {cleaned_utterance}")
+                            last_role = current_role # Update last role
+                    # Apply max turns limit to the formatted turns
+                    formatted_dialogue = formatted_dialogue[:CONFIG["max_turns"]]
+                    if formatted_dialogue:
+                        text_samples.append(" </s> ".join(formatted_dialogue))
+                        processed_conv_count += 1 # Count processed unique conversations
+
+            except Exception as e:
+                logging.error(f"Failed to load or process empathetic_dialogues for tokenizer: {e}")
+
+
+        # --- Process BlendedSkillTalk ---
+        if "blended_skill_talk" in dataset_names:
+            logging.info(f"Loading blended_skill_talk for tokenizer training (max {samples_per_dataset} dialogues)...")
+            try:
+                # Load dialogues - BST is structured differently, slice directly
+                bst_dataset = load_dataset("blended_skill_talk", split=f"train[:{samples_per_dataset}]", trust_remote_code=True) # Add trust_remote_code=True
+                logging.info("Processing blended_skill_talk...")
+                for entry in bst_dataset:
+                    formatted_dialogue = []
+                    # Combine the dialogue history and the final two turns
+                    dialogue_turns_raw = entry['previous_utterance']
+                    # Add final utterances if they exist and are not empty strings
+                    if entry.get('free_turker_utterance'):
+                        dialogue_turns_raw.append(entry['free_turker_utterance'])
+                    if entry.get('guided_turker_utterance'):
+                         dialogue_turns_raw.append(entry['guided_turker_utterance'])
+
+                    turns_to_process = dialogue_turns_raw[:CONFIG["max_turns"]] # Apply max turns limit
+                    for i, utterance in enumerate(turns_to_process):
+                        role = "<user>" if i % 2 == 0 else "<assistant>" # Assume simple alternation
+                        cleaned_utterance = self._clean_text(utterance)
+                        if cleaned_utterance:
+                            formatted_dialogue.append(f"{role} {cleaned_utterance}")
+                    if formatted_dialogue:
+                        text_samples.append(" </s> ".join(formatted_dialogue))
+            except Exception as e:
+                logging.error(f"Failed to load or process blended_skill_talk for tokenizer: {e}")
+
+        # --- Process PersonaChat ---
+        if "AlekseyKorshuk/persona-chat" in dataset_names: # Correct dataset identifier
+            pc_dataset_name = "AlekseyKorshuk/persona-chat"
+            logging.info(f"Loading {pc_dataset_name} for tokenizer training (max {samples_per_dataset} dialogues)...")
+            try:
+                pc_dataset = load_dataset(pc_dataset_name, split=f"train[:{samples_per_dataset}]", trust_remote_code=True) # Add trust_remote_code=True, Correct dataset identifier
+                logging.info(f"Processing {pc_dataset_name}...")
+                for entry in pc_dataset:
+                    # PersonaChat often has 'utterances' containing 'history'
+                    if 'utterances' in entry and entry['utterances']:
+                        # Get the history from the last item in utterances for the full dialogue
+                        history = entry['utterances'][-1]['history']
+                        history = history[:CONFIG["max_turns"]] # Apply max turns
+                        formatted_dialogue = []
+                        for i, utterance in enumerate(history):
+                             role = "<user>" if i % 2 == 0 else "<assistant>" # Assume simple alternation
+                             cleaned_utterance = self._clean_text(utterance)
+                             if cleaned_utterance:
+                                  formatted_dialogue.append(f"{role} {cleaned_utterance}")
+                        if formatted_dialogue:
+                            text_samples.append(" </s> ".join(formatted_dialogue))
+                    else:
+                        logging.warning(f"Skipping {pc_dataset_name} entry due to unexpected structure: {entry}")
+
+            except Exception as e:
+                logging.error(f"Failed to load or process {pc_dataset_name} for tokenizer: {e}")
+
+        # --- Process Universal-Transformers-RawConvo ---
+        if "elapt1c/Universal-Transformers-RawConvo" in dataset_names:
+            utrc_dataset_name = "elapt1c/Universal-Transformers-RawConvo"
+            logging.info(f"Loading {utrc_dataset_name} for tokenizer training (max {samples_per_dataset} dialogues)...")
+            try:
+                utrc_dataset = load_dataset(utrc_dataset_name, split=f"train[:{samples_per_dataset}]", trust_remote_code=True)
+                logging.info(f"Processing {utrc_dataset_name}...")
+                for entry in utrc_dataset:
+                    formatted_dialogue = []
+                    dialogue = entry['dialogue'] # Assuming 'dialogue' column contains the conversation turns
+                    if isinstance(dialogue, list): # Check if dialogue is a list of strings
+                        dialogue = dialogue[:CONFIG["max_turns"]]
+                        for i, utterance in enumerate(dialogue):
+                            role = "<user>" if i % 2 == 0 else "<assistant>"
+                            cleaned_utterance = self._clean_text(utterance)
+                            if cleaned_utterance:
+                                formatted_dialogue.append(f"{role} {cleaned_utterance}")
+                        if formatted_dialogue:
+                            text_samples.append(" </s> ".join(formatted_dialogue))
+                    else:
+                        logging.warning(f"Skipping {utrc_dataset_name} entry due to unexpected 'dialogue' structure: {type(dialogue)}")
+
+            except Exception as e:
+                logging.error(f"Failed to load or process {utrc_dataset_name} for tokenizer: {e}")
+
+        # --- Process papahawk/conversational-01 ---
+        if "papahawk/conversational-01" in dataset_names:
+            pco_dataset_name = "papahawk/conversational-01"
+            logging.info(f"Loading {pco_dataset_name} for tokenizer training (max {samples_per_dataset} samples)...")
+            try:
+                pco_dataset = load_dataset(pco_dataset_name, split=f"train[:{samples_per_dataset}]", trust_remote_code=True)
+                logging.info(f"Processing {pco_dataset_name}...")
+                for entry in pco_dataset:
+                    prompt = entry['prompt']
+                    response = entry['response']
+                    cleaned_prompt = self._clean_text(prompt)
+                    cleaned_response = self._clean_text(response)
+                    if cleaned_prompt and cleaned_response:
+                        formatted_dialogue = [f"<user> {cleaned_prompt}", f"<assistant> {cleaned_response}"]
+                        text_samples.append(" </s> ".join(formatted_dialogue))
+
+            except Exception as e:
+                logging.error(f"Failed to load or process {pco_dataset_name} for tokenizer: {e}")
+
+
+        logging.info(f"Total text samples for tokenizer training: {len(text_samples)}")
+        if not text_samples:
+            raise ValueError("No text samples collected for tokenizer training. Check dataset loading and paths.")
+
+        # Ensure tokenizer directory exists before training
+        os.makedirs(self.tokenizer_dir, exist_ok=True)
+
+        logging.info(f"Training BPE tokenizer with vocab size {CONFIG['vocab_size']}...")
+        trainer = trainers.BpeTrainer(
+            vocab_size=CONFIG["vocab_size"],
+            special_tokens=self.special_tokens,
+            min_frequency=2, # Keep min_frequency low with more data
+            show_progress=True
+        )
+        # Make sure text_samples is an iterator or list of strings
+        def text_iterator():
+            for sample in text_samples:
+                yield sample
+
+        self.tokenizer.train_from_iterator(text_iterator(), trainer=trainer, length=len(text_samples))
+
+        eos_token_id = self.tokenizer.token_to_id("</s>")
+        if eos_token_id is None:
+            logging.warning("</s> token not found in trained tokenizer vocab! Using <pad> as fallback for post-processor.")
+            eos_token_id = self.tokenizer.token_to_id("<pad>") or 0 # Fallback needed
+
+        # Configure post-processor (adjust if needed based on how you structure input/output)
+        self.tokenizer.post_processor = processors.TemplateProcessing(
+            single="$A </s>",
+            pair="$A </s> $B </s>", # How to handle pairs - maybe just use single always?
+            special_tokens=[("</s>", eos_token_id)],
+        )
+
+        logging.info(f"Saving tokenizer to {self.tokenizer_path}")
+        self.tokenizer.save(self.tokenizer_path)
+        logging.info("Tokenizer training complete.")
+
+    def get_tokenizer(self):
+         if not os.path.exists(self.tokenizer_path):
+              raise FileNotFoundError(f"Tokenizer file not found at {self.tokenizer_path}. Train tokenizer first.")
+         tokenizer = Tokenizer.from_file(self.tokenizer_path)
+         # Verify special tokens crucial for processing exist
+         required_tokens = ["<pad>", "<s>", "</s>", "<unk>", "<user>", "<assistant>"]
+         for token in required_tokens:
+              if tokenizer.token_to_id(token) is None:
+                   raise ValueError(f"Crucial special token '{token}' not found in loaded tokenizer '{self.tokenizer_path}'!")
+         return tokenizer
+
+# --- Dataset Loading and Processing ---
+
+class CombinedChatDataset(Dataset):
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.pad_id = self.tokenizer.token_to_id("<pad>")
+        self.eos_id = self.tokenizer.token_to_id("</s>")
+        self.bos_id = self.tokenizer.token_to_id("<s>")
+        self.user_id = self.tokenizer.token_to_id("<user>")
+        self.assistant_id = self.tokenizer.token_to_id("<assistant>")
+        self.max_length = CONFIG["max_seq_len"]
+        # Reuse cleaning function from TokenizerTrainer instance
+        self._clean_text = TokenizerTrainer()._clean_text
+
+        self.all_processed_conversations = []
+
+        # --- Process DailyDialog ---
+        if "daily_dialog" in CONFIG["datasets"]:
+            logging.info("Loading and processing daily_dialog dataset...")
+            try:
+                dd_dataset = load_dataset("daily_dialog", split="train", trust_remote_code=True) # Add trust_remote_code=True
+                logging.info(f"Processing {len(dd_dataset)} daily_dialog conversations...")
+                for entry in dd_dataset:
+                    conversation = []
+                    dialogue = entry['dialog'][:CONFIG["max_turns"]]
+                    if not dialogue: continue
+                    for i, utterance in enumerate(dialogue):
+                        role = "<user>" if i % 2 == 0 else "<assistant>"
+                        cleaned_text = self._clean_text(utterance)
+                        if cleaned_text:
+                            conversation.append({'role': role, 'text': cleaned_text})
+                    if conversation:
+                        self.all_processed_conversations.append(conversation)
+            except Exception as e:
+                 logging.error(f"Failed to load or process daily_dialog for training: {e}")
+
+        # --- Process EmpatheticDialogues ---
+        if "empathetic_dialogues" in CONFIG["datasets"]:
+            logging.info("Loading and processing empathetic_dialogues dataset...")
+            try:
+                ed_dataset = load_dataset("empathetic_dialogues", split="train", trust_remote_code=True) # Add trust_remote_code=True
+                logging.info("Grouping empathetic_dialogues by conversation ID...")
+                conversations_grouped = defaultdict(list)
+                for entry in ed_dataset:
+                    conversations_grouped[entry['conv_id']].append(entry)
+
+                logging.info(f"Processing {len(conversations_grouped)} empathetic_dialogues conversations...")
+                for conv_id, entries in conversations_grouped.items():
+                    conversation = []
+                    sorted_entries = sorted(entries, key=lambda x: x['utterance_idx'])
+                    # Handle context as first user turn if present
+                    if sorted_entries[0]['context']:
+                        context_text = self._clean_text(sorted_entries[0]['context'])
+                        if context_text:
+                             conversation.append({'role': '<user>', 'text': context_text})
+                    # Process utterances, assuming alternation
+                    last_role = conversation[-1]['role'] if conversation else None # Role of the last added turn
+                    for entry in sorted_entries:
+                         text = self._clean_text(entry['utterance'])
+                         if not text: continue
+                         # Determine role based on the *last added* role
+                         current_role = '<assistant>' if last_role == '<user>' else '<user>'
+                         conversation.append({'role': current_role, 'text': text})
+                         last_role = current_role # Update for next iteration
+
+                    # Apply max turns limit *after* forming the full sequence
+                    conversation = conversation[:CONFIG["max_turns"]]
+                    if conversation:
+                        self.all_processed_conversations.append(conversation)
+
+            except Exception as e:
+                logging.error(f"Failed to load or process empathetic_dialogues for training: {e}")
+
+        # --- Process BlendedSkillTalk ---
+        if "blended_skill_talk" in CONFIG["datasets"]:
+            logging.info("Loading and processing blended_skill_talk dataset...")
+            try:
+                bst_dataset = load_dataset("blended_skill_talk", split="train", trust_remote_code=True) # Add trust_remote_code=True
+                logging.info(f"Processing {len(bst_dataset)} blended_skill_talk conversations...")
+                for entry in bst_dataset:
+                    conversation = []
+                    # Reconstruct dialogue: history + final two turns (if they exist)
+                    dialogue_turns_raw = entry['previous_utterance']
+                    if entry.get('free_turker_utterance'):
+                        dialogue_turns_raw.append(entry['free_turker_utterance'])
+                    if entry.get('guided_turker_utterance'):
+                         dialogue_turns_raw.append(entry['guided_turker_utterance'])
+
+                    if not dialogue_turns_raw: continue # Skip if no turns found
+
+                    turns_to_process = dialogue_turns_raw[:CONFIG["max_turns"]] # Apply max turns limit
+
+                    for i, utterance in enumerate(turns_to_process):
+                        role = "<user>" if i % 2 == 0 else "<assistant>" # Assume simple alternation
+                        cleaned_text = self._clean_text(utterance)
+                        if cleaned_text:
+                            conversation.append({'role': role, 'text': cleaned_text})
+                    if conversation: # Only add if not empty after cleaning/truncation
+                        self.all_processed_conversations.append(conversation)
+            except Exception as e:
+                logging.error(f"Failed to load or process blended_skill_talk for training: {e}")
+
+        # --- Process PersonaChat ---
+        if "AlekseyKorshuk/persona-chat" in CONFIG["datasets"]: # Correct dataset identifier
+            pc_dataset_name = "AlekseyKorshuk/persona-chat"
+            logging.info(f"Loading and processing {pc_dataset_name} dataset...")
+            try:
+                pc_dataset = load_dataset(pc_dataset_name, split="train", trust_remote_code=True) # Add trust_remote_code=True, Correct dataset identifier
+                logging.info(f"Processing {len(pc_dataset)} {pc_dataset_name} conversations...")
+                for entry in pc_dataset:
+                    conversation = []
+                    if 'utterances' in entry and entry['utterances']:
+                        # Extract the dialogue history
+                        history = entry['utterances'][-1]['history']
+                        history = history[:CONFIG["max_turns"]] # Apply max turns limit
+
+                        for i, utterance in enumerate(history):
+                            role = "<user>" if i % 2 == 0 else "<assistant>" # Simple alternation
+                            cleaned_text = self._clean_text(utterance)
+                            if cleaned_text:
+                                conversation.append({'role': role, 'text': cleaned_text})
+
+                        if conversation: # Only add if not empty
+                            self.all_processed_conversations.append(conversation)
+                    else:
+                         logging.warning(f"Skipping {pc_dataset_name} entry due to unexpected structure: {entry.keys()}")
+
+            except Exception as e:
+                logging.error(f"Failed to load or process {pc_dataset_name} for training: {e}")
+
+        # --- Process Universal-Transformers-RawConvo ---
+        if "elapt1c/Universal-Transformers-RawConvo" in CONFIG["datasets"]:
+            utrc_dataset_name = "elapt1c/Universal-Transformers-RawConvo"
+            logging.info(f"Loading and processing {utrc_dataset_name} dataset...")
+            try:
+                utrc_dataset = load_dataset(utrc_dataset_name, split="train", trust_remote_code=True)
+                logging.info(f"Processing {len(utrc_dataset)} {utrc_dataset_name} conversations...")
+                for entry in utrc_dataset:
+                    conversation = []
+                    dialogue = entry['dialogue'] # Assuming 'dialogue' column contains the conversation turns
+                    if isinstance(dialogue, list): # Check if dialogue is a list of strings
+                        dialogue = dialogue[:CONFIG["max_turns"]]
+                        for i, utterance in enumerate(dialogue):
+                            role = "<user>" if i % 2 == 0 else "<assistant>"
+                            cleaned_text = self._clean_text(utterance)
+                            if cleaned_text:
+                                conversation.append({'role': role, 'text': cleaned_text})
+                        if conversation:
+                            self.all_processed_conversations.append(conversation)
+                    else:
+                        logging.warning(f"Skipping {utrc_dataset_name} entry due to unexpected 'dialogue' structure: {type(dialogue)}")
+
+            except Exception as e:
+                logging.error(f"Failed to load or process {utrc_dataset_name} for training: {e}")
+
+        # --- Process papahawk/conversational-01 ---
+        if "papahawk/conversational-01" in CONFIG["datasets"]:
+            pco_dataset_name = "papahawk/conversational-01"
+            logging.info(f"Loading and processing {pco_dataset_name} dataset...")
+            try:
+                pco_dataset = load_dataset(pco_dataset_name, split="train", trust_remote_code=True)
+                logging.info(f"Processing {len(pco_dataset)} {pco_dataset_name} conversations...")
+                for entry in pco_dataset:
+                    prompt = entry['prompt']
+                    response = entry['response']
+                    cleaned_prompt = self._clean_text(prompt)
+                    cleaned_response = self._clean_text(response)
+                    if cleaned_prompt and cleaned_response:
+                        conversation = [
+                            {'role': '<user>', 'text': cleaned_prompt},
+                            {'role': '<assistant>', 'text': cleaned_response}
+                        ]
+                        self.all_processed_conversations.append(conversation)
+
+            except Exception as e:
+                logging.error(f"Failed to load or process {pco_dataset_name} for training: {e}")
+
+
+        logging.info(f"Total processed conversations from all datasets: {len(self.all_processed_conversations)}")
+        if not self.all_processed_conversations:
+             raise ValueError("No processed conversations were created from any dataset. Check loading logic and dataset availability.")
+
+        logging.info("Shuffling combined dataset...")
+        random.shuffle(self.all_processed_conversations)
+
+
+    def __len__(self):
+        return len(self.all_processed_conversations)
+
+    def __getitem__(self, idx):
+        conversation = self.all_processed_conversations[idx]
+        formatted_ids = [self.bos_id]
+        for turn in conversation:
+            role_id = self.user_id if turn['role'] == '<user>' else self.assistant_id
+            # Encode without adding special tokens automatically by tokenizer
+            try:
+                utterance_ids = self.tokenizer.encode(turn['text'], add_special_tokens=False).ids
+            except Exception as e:
+                 logging.error(f"Error encoding text at index {idx}, turn '{turn}': {e}")
+                 utterance_ids = [] # Skip this utterance on error
+
+            # Check length: Current + Role + Utterance + EOS <= MaxLength
+            # Need +1 for role, +len(utterance), +1 for potential EOS
+            if len(formatted_ids) + 1 + len(utterance_ids) + 1 > self.max_length:
+                # Attempt to add just the role and EOS if utterance is too long
+                if len(formatted_ids) + 1 + 1 <= self.max_length:
+                     formatted_ids.append(role_id)
+                     formatted_ids.append(self.eos_id)
+                break # Stop adding turns
+
+            formatted_ids.append(role_id)
+            formatted_ids.extend(utterance_ids)
+            formatted_ids.append(self.eos_id)
+
+        # Final safety truncate (should be rare if logic above is correct)
+        if len(formatted_ids) > self.max_length:
+             formatted_ids = formatted_ids[:self.max_length]
+             # Ensure last token isn't partial (though unlikely with BPE)
+             # If the truncated sequence ends with a role ID, it's probably bad, remove it.
+             if formatted_ids and (formatted_ids[-1] == self.user_id or formatted_ids[-1] == self.assistant_id):
+                  formatted_ids.pop()
+             # If after popping the role ID, it's still too long (unlikely), truncate again
+             if len(formatted_ids) > self.max_length:
+                 formatted_ids = formatted_ids[:self.max_length]
+
+
+        # Handle case of extremely short sequences after processing
+        if len(formatted_ids) < 2: # Need at least BOS and one other token for input/label pair
+             logging.warning(f"Sequence at index {idx} is too short after processing (<2 tokens). Skipping. Original length: {len(conversation)}")
+             # Return None to be filtered by collate_fn
+             return None
+
+        input_ids = formatted_ids[:-1]
+        labels = formatted_ids[1:]
+
+        # Final check before returning
+        if len(input_ids) == 0:
+            logging.warning(f"Sequence at index {idx} resulted in empty input_ids after slicing. Skipping.")
+            return None
+
+
+        return {"input_ids": input_ids, "labels": labels}
+
+    @staticmethod
+    def collate_fn(batch):
+        # Filter out None items from __getitem__
+        batch = [item for item in batch if item is not None]
+        if not batch:
+            return None # Return None if the whole batch was invalid
+
+        max_len = max(len(item["input_ids"]) for item in batch)
+
+        # Load tokenizer once to get pad_id - ensure path matches CONFIG
+        try:
+            # Correctly reference the tokenizer path from CONFIG within the static method
+            tokenizer_path = os.path.join("tokenizer", CONFIG["tokenizer_name"])
+            # TODO: Consider passing tokenizer/pad_id if this becomes a bottleneck
+            tokenizer = Tokenizer.from_file(tokenizer_path)
+            pad_id = tokenizer.token_to_id("<pad>")
+            if pad_id is None: raise ValueError("<pad> token not found")
+        except Exception as e:
+            logging.error(f"Collate Error: Failed to load tokenizer or get pad_id ('{CONFIG['tokenizer_name']}'): {e}")
+            pad_id = 0 # Risky fallback
+
+        inputs, labels, masks = [], [], []
+        for item in batch:
+            input_len = len(item["input_ids"])
+            pad_len = max_len - input_len
+            inputs.append(item["input_ids"] + [pad_id] * pad_len)
+            # Pad labels with pad_id (or any ID to be ignored by CrossEntropyLoss)
+            labels.append(item["labels"] + [pad_id] * pad_len)
+            masks.append([1] * input_len + [0] * pad_len)
+
+        return {
+            "input_ids": torch.tensor(inputs, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+            "attention_mask": torch.tensor(masks, dtype=torch.long) # Or bool
+        }
+
+# --- Trainer, Safety Manager, Checkpoint Manager ---
+
+class HROMTrainer:
+    def __init__(self, model, tokenizer):
+        # Determine device based on number of GPUs
+        if CONFIG["num_gpus"] > 1 and torch.cuda.device_count() >= CONFIG["num_gpus"]:
+            self.device = [torch.device(f"cuda:{i}") for i in range(CONFIG["num_gpus"])]
+            logging.info(f"Using DataParallel on devices: {self.device}")
+            model = nn.DataParallel(model, device_ids=self.device) # Wrap model for DP
+            self.main_device = self.device[0] # Main device for logging etc.
+        elif torch.cuda.is_available():
+            self.device = [torch.device("cuda:0")] # Single GPU list for consistency
+            self.main_device = self.device[0]
+            logging.info(f"Using single GPU: {self.main_device}")
+        else:
+            self.device = [torch.device("cpu")]
+            self.main_device = self.device[0]
+            logging.info(f"Using CPU: {self.main_device}")
+
+        self.model = model.to(self.main_device) # Move model to main device
+
+        self.use_amp = (self.main_device.type == "cuda" and hasattr(torch.cuda.amp, "GradScaler"))
+        self.scaler = torch.cuda.amp.GradScaler() if self.use_amp else None
+        logging.info(f"Automatic Mixed Precision (AMP): {'Enabled' if self.use_amp else 'Disabled'}")
+
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=CONFIG["learning_rate"], # Base LR
+            betas=(0.9, 0.95),
+            weight_decay=0.1,
+            fused= (self.main_device.type == "cuda")
+        )
+        self.tokenizer = tokenizer
+        self.pad_id = self.tokenizer.token_to_id("<pad>")
+        if self.pad_id is None:
+             # Attempt to get from config if available or fallback
+             self.pad_id = CONFIG.get("pad_token_id", 0)
+             logging.warning(f"<pad> token ID not found in tokenizer, using fallback ID: {self.pad_id}")
+
+
+        # Make sure ignore_index uses the determined pad_id
+        self.criterion = nn.CrossEntropyLoss(ignore_index=self.pad_id)
+        self.base_lr = CONFIG["learning_rate"]
+        self.warmup_steps = CONFIG["warmup_steps"]
+
+    def _adjust_learning_rate(self, step):
+        if self.warmup_steps > 0 and step < self.warmup_steps:
+            lr = self.base_lr * (step + 1) / self.warmup_steps
+        else:
+            # Optional: Add LR decay (e.g., cosine) after warmup
+            # Example: lr = self.base_lr * 0.5 * (1 + math.cos(math.pi * (step - self.warmup_steps) / (total_steps - self.warmup_steps)))
+            lr = self.base_lr # Keep base LR after warmup for now
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
+        return lr
+
+    def train_step(self, batch):
+        # Determine precision for autocast
+        if self.use_amp:
+            # Use torch.float16 (FP16) for mixed precision training
+            amp_dtype = torch.float16 # FP16 training explicitly requested
+        autocast_context = torch.cuda.amp.autocast(dtype=amp_dtype, enabled=self.use_amp) if self.use_amp else nullcontext()
+
+        with autocast_context:
+            input_ids = batch["input_ids"].to(self.main_device) # Move to main device
+            attention_mask = batch["attention_mask"].to(self.main_device) # Move to main device
+            labels = batch["labels"].to(self.main_device) # Move to main device
+
+            outputs = self.model(input_ids, attention_mask=attention_mask)
+
+            # Reshape for loss calculation
+            logits_flat = outputs.view(-1, outputs.size(-1)) # Shape: (B * T, vocab_size)
+            labels_flat = labels.view(-1)                   # Shape: (B * T)
+
+            # Calculate loss - ensure logits are float32 for stability esp. with AMP
+            loss = self.criterion(logits_flat.float(), labels_flat)
+
+            # Scale loss for gradient accumulation
+            scaled_loss = loss / CONFIG["grad_accum_steps"]
+
+        # Backward pass
+        if self.use_amp and self.scaler:
+            self.scaler.scale(scaled_loss).backward()
+        else:
+            scaled_loss.backward()
+
+        return loss.item() # Return the unscaled loss for logging
+
+    def clip_and_step(self, current_optimizer_step):
+         current_lr = self._adjust_learning_rate(current_optimizer_step)
+         # Gradient Clipping *before* optimizer step
+         if self.use_amp and self.scaler:
+             # Unscale first - important before clipping
+             self.scaler.unscale_(self.optimizer)
+             # Clip grad norm
+             torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+             # Optimizer step (with scaler)
+             self.scaler.step(self.optimizer)
+             # Update scaler for next iteration
+             self.scaler.update()
+         else:
+             # Clip grad norm
+             torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+             # Optimizer step
+             self.optimizer.step()
+
+         # Zero gradients *after* stepping
+         self.optimizer.zero_grad(set_to_none=True)
+         return current_lr
+
+
+class SafetyManager:
+    def __init__(self, model, tokenizer, main_device): # Added main_device argument
+        self.model = model
+        self.tokenizer = tokenizer
+        self.main_device = main_device # Store main_device
+        # More conservative list
+        self.bad_words = ["kill", "murder", "suicide", "hate", "abuse", "violence", "illegal", "harm", "die", "attack", "rape", "molest", "exploit", "terror"]
+        self.bad_word_ids = []
+        logging.info("Initializing safety manager...")
+        # Pre-encode bad word sequences
+        for word in self.bad_words:
+             # Encode potentially multi-token words carefully
+             ids = tokenizer.encode(f" {word}", add_special_tokens=False).ids # Add prefix space for BPE
+             if ids:
+                self.bad_word_ids.append(ids)
+                logging.debug(f"Encoded bad word '{word}' (with space) to IDs: {ids}")
+             # Try without space too
+             ids_no_space = tokenizer.encode(word, add_special_tokens=False).ids
+             if ids_no_space and ids_no_space != ids:
+                  self.bad_word_ids.append(ids_no_space)
+                  logging.debug(f"Encoded bad word '{word}' (no space) to IDs: {ids_no_space}")
+
+             if not ids and not ids_no_space:
+                logging.warning(f"Could not encode bad word '{word}' - skipping.")
+
+        # Pre-get special IDs
+        self.eos_id = self.tokenizer.token_to_id("</s>")
+        self.bos_id = self.tokenizer.token_to_id("<s>")
+        self.user_id = self.tokenizer.token_to_id("<user>")
+        self.assistant_id = self.tokenizer.token_to_id("<assistant>")
+        self.pad_id = self.tokenizer.token_to_id("<pad>")
+
+        if self.eos_id is None: logging.error("</s> token ID not found for SafetyManager!"); self.eos_id = 0
+        if self.bos_id is None: logging.error("<s> token ID not found for SafetyManager!"); self.bos_id = 0
+        if self.user_id is None: logging.error("<user> token ID not found for SafetyManager!")
+        if self.assistant_id is None: logging.error("<assistant> token ID not found for SafetyManager!")
+        if self.pad_id is None: logging.error("<pad> token ID not found for SafetyManager!"); self.pad_id = 0
+
+
+    def contains_sequence(self, tokens, seq):
+        """Checks if the list `tokens` contains the sublist `seq`."""
+        if not seq or not tokens or len(tokens) < len(seq):
+            return False
+        seq_len = len(seq)
+        for i in range(len(tokens) - seq_len + 1):
+            if tokens[i : i + seq_len] == seq:
+                return True
+        return False
+
+    def content_filter(self, text_ids):
+        """Checks if a list of token IDs contains any bad word sequences."""
+        if not isinstance(text_ids, list):
+            logging.warning("Content filter received non-list input.")
+            return True # Default to safe if input is weird
+        for bad_ids in self.bad_word_ids:
+            if self.contains_sequence(text_ids, bad_ids):
+                # Log the detected sequence for debugging
+                detected_word = self.tokenizer.decode(bad_ids)
+                logging.warning(f"Unsafe content detected: Found sequence corresponding to '{detected_word}' (IDs: {bad_ids}).")
+                return False # Unsafe
+        return True # Safe
+
+    def generate_safely(self, prompt, max_new_tokens=50, temperature=0.5, top_k=50):
+        self.model.eval()
+        device = self.main_device # Use stored main_device for generation
+        if isinstance(self.main_device, list): # Handle DataParallel case for generation device - SHOULD BE self.main_device now
+            device = self.main_device[0] # Ensure generation happens on the main GPU
+
+
+        # Encode prompt, ensure it ends appropriately (e.g., with role token + EOS?)
+        # Let's assume the prompt ends like "<user> blah blah </s>" and we need to add "<assistant>"
+        prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False).ids
+
+        # Start generation sequence with BOS, prompt, and assistant token
+        # Ensure prompt doesn't already include BOS
+        if prompt_ids and prompt_ids[0] == self.bos_id:
+             input_ids = list(prompt_ids)
+        else:
+             input_ids = [self.bos_id] + list(prompt_ids)
+
+        # Add the assistant token to signal the model to generate the response
+        if self.assistant_id is not None:
+            input_ids.append(self.assistant_id)
+        else:
+            logging.error("Assistant token ID is None, cannot properly start generation.")
+            return "Error: Assistant token not found."
+
+
+        generated_ids = list(input_ids) # Start with the prepared input sequence
+        logging.debug(f"Starting safe generation with initial IDs: {generated_ids}")
+
+        with torch.no_grad():
+            for step in range(max_new_tokens):
+                # Prepare input tensor for this step - only use up to max_seq_len
+                current_input_ids = generated_ids[-CONFIG["max_seq_len"]:]
+                current_input_tensor = torch.tensor([current_input_ids]).to(device)
+                # Create attention mask for the current length
+                attention_mask = torch.ones_like(current_input_tensor)
+
+                # Model forward pass
+                try:
+                    outputs = self.model(current_input_tensor, attention_mask=attention_mask)
+                    next_token_logits = outputs[:, -1, :] # Logits for the next token
+                    if isinstance(self.model, nn.DataParallel): # Handle DP output
+                        next_token_logits = next_token_logits.mean(dim=0, keepdim=True) # Average logits from DP
+                except Exception as e:
+                     logging.error(f"Model forward pass failed during generation: {e}")
+                     break # Stop generation on error
+
+                # --- Safety Check BEFORE sampling ---
+                # Apply penalties to bad word starting tokens if possible
+                # For now, we filter *after* sampling the token
+
+                # Sampling (Temperature, Top-K)
+                if temperature > 0 and temperature != 1.0:
+                    next_token_logits = next_token_logits / temperature
+                if top_k > 0 and top_k < next_token_logits.size(-1): # Ensure top_k is valid
+                    v, _ = torch.topk(next_token_logits, top_k)
+                    # Handle potential NaN/Inf in logits before comparison
+                    safe_logits = torch.nan_to_num(next_token_logits, nan=-float('inf'), posinf=float('inf'), neginf=-float('inf'))
+                    threshold = v[:, [-1]]
+                    safe_logits[safe_logits < threshold] = -float('Inf')
+                    next_token_logits = safe_logits # Use the filtered logits
+
+                probs = torch.softmax(next_token_logits, dim=-1)
+                # Handle potential NaNs in probabilities before sampling
+                if torch.isnan(probs).any():
+                     logging.warning("NaN detected in probabilities before sampling. Replacing with uniform distribution.")
+                     probs = torch.ones_like(probs) / probs.size(-1) # Fallback to uniform
+
+                next_token_id = torch.multinomial(probs, num_samples=1).item()
+
+                # --- Safety Check AFTER sampling token ---
+                # Check if adding this token creates a bad sequence
+                potential_sequence_ids = generated_ids + [next_token_id]
+                # Check only the newly formed part for bad words for efficiency?
+                # Let's check the whole sequence for simplicity/robustness for now.
+                if not self.content_filter(potential_sequence_ids):
+                    logging.warning(f"Potential unsafe token ({next_token_id}, '{self.tokenizer.decode([next_token_id])}') blocked POST-sampling. Stopping generation.")
+                    # Optionally try sampling a different token? For now, just stop.
+                    break
+
+                # Add the safe token
+                generated_ids.append(next_token_id)
+
+                # Check for EOS token
+                if next_token_id == self.eos_id:
+                    logging.debug(f"EOS token generated at step {step+1}. Stopping generation.")
+                    break
+
+                # Prevent infinite loops if max tokens reached
+                if step == max_new_tokens - 1:
+                     logging.debug("Max new tokens reached. Stopping generation.")
+                     # Ensure the sequence ends with EOS if it didn't naturally
+                     if generated_ids[-1] != self.eos_id and self.eos_id is not None:
+                         generated_ids.append(self.eos_id)
+
+        self.model.train() # Set model back to training mode
+
+        # Decode the generated part (excluding the initial prompt + assistant token)
+        start_index = len(input_ids)
+        response_ids = generated_ids[start_index:]
+
+        # Decode, skipping special tokens like EOS, BOS, PAD but potentially keeping USER/ASSISTANT
+        # Let's skip all special tokens for the final output text for clarity.
+        decoded_text = self.tokenizer.decode(response_ids, skip_special_tokens=True).strip()
+
+        return decoded_text
+
+
+    def debug_generation(self, prompt="<user> Tell me about your hobbies."): # Example prompt
+         logging.info(f"\n--- Debug Generation & Safety Check ---")
+         # Ensure prompt ends logically for the model (e.g., with user token and EOS)
+         if not prompt.strip().endswith("</s>"):
+              if not prompt.strip().endswith("<user>") and not prompt.strip().endswith("<assistant>"):
+                   prompt = prompt.strip() + " </s>" # Add EOS if ends mid-sentence
+              else:
+                   prompt = prompt.strip() + " </s>" # Add EOS after role token
+
+         # Ensure the prompt starts appropriately (e.g., no BOS needed here as generate_safely adds it)
+         if prompt.startswith("<s>"):
+              prompt = prompt[len("<s>"):].strip()
+
+
+         generated_response = self.generate_safely(prompt, max_new_tokens=60, temperature=0.7, top_k=50)
+
+         logging.info(f"Prompt Sent: '{prompt}'")
+         logging.info(f"Generated Response: '{generated_response}'")
+         logging.info("\n--- End Debug Generation ---\n")
+
+
+class CheckpointManager:
+    def __init__(self):
+        # Use checkpoint directory from CONFIG
+        self.checkpoint_dir = CONFIG["checkpoint_dir"]
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+        logging.info(f"Checkpoint directory set to: {self.checkpoint_dir}")
+
+    def save(self, model, optimizer, step):
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        # Use a consistent naming scheme based on the directory name if desired
+        prefix = os.path.basename(self.checkpoint_dir).replace("checkpoints_", "")
+        # Ensure step is converted to string if it's passed as something else (e.g., 'final')
+        step_str = str(step)
+        filename = f"hrom_{prefix}_step{step_str}_{timestamp}.pt"
+        path = os.path.join(self.checkpoint_dir, filename)
+        state = {
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "step": step if isinstance(step, int) else -1, # Store step number or -1 for non-numeric steps
+            "config": CONFIG # Save config with checkpoint
+        }
+        logging.info(f"Saving checkpoint to {path}...")
+        try:
+            torch.save(state, path)
+            logging.info(f"Checkpoint saved successfully at step {step_str}.")
+            self._cleanup_old_checkpoints()
+        except Exception as e:
+            logging.error(f"Failed to save checkpoint '{path}': {e}")
+
+    def _cleanup_old_checkpoints(self):
+        max_checkpoints = CONFIG.get("max_checkpoints", 5) # Get from config, default 5
+        if max_checkpoints <= 0:
+             return # Keep all checkpoints if max_checkpoints is non-positive
+
+        try:
+            # Filter only files matching the expected pattern (avoid deleting other files)
+            prefix = os.path.basename(self.checkpoint_dir).replace("checkpoints_", "")
+            pattern = re.compile(rf"hrom_{prefix}_step(\d+|.+)_(\d{{8}}_\d{{6}})\.pt")
+
+            checkpoints = []
+            for f in os.listdir(self.checkpoint_dir):
+                 match = pattern.match(f)
+                 if match:
+                      filepath = os.path.join(self.checkpoint_dir, f)
+                      checkpoints.append((filepath, os.path.getmtime(filepath)))
+
+            # Sort by modification time (oldest first)
+            checkpoints.sort(key=lambda x: x[1])
+
+            num_to_delete = len(checkpoints) - max_checkpoints
+            if num_to_delete > 0:
+                #logging.info(f"Max checkpoints ({max_checkpoints}) reached. Removing {num_to_delete} oldest checkpoints.")
+                for i in range(num_to_delete):
+                    file_to_remove, _ = checkpoints[i]
+                    try:
+                        os.remove(file_to_remove)
+                        #logging.info(f"Removed old checkpoint: {os.path.basename(file_to_remove)}")
+                    except OSError as e:
+                        logging.error(f"Error removing checkpoint {file_to_remove}: {e}")
+        except Exception as e:
+            logging.error(f"Error during checkpoint cleanup: {e}")
+
+
+    def load_latest(self, model, optimizer):
+        try:
+            # Filter files based on pattern and sort by time
+            prefix = os.path.basename(self.checkpoint_dir).replace("checkpoints_", "")
+            pattern = re.compile(rf"hrom_{prefix}_step(\d+|.+)_(\d{{8}}_\d{{6}})\.pt")
+            checkpoints = []
+            for f in os.listdir(self.checkpoint_dir):
+                 match = pattern.match(f)
+                 if match:
+                      filepath = os.path.join(self.checkpoint_dir, f)
+                      checkpoints.append((filepath, os.path.getmtime(filepath)))
+
+            if not checkpoints:
+                logging.info("No valid checkpoints found to load.")
+                return 0 # Start from step 0
+
+            # Sort by modification time (newest first)
+            checkpoints.sort(key=lambda x: x[1], reverse=True)
+
+            latest_checkpoint_path, _ = checkpoints[0]
+            logging.info(f"Loading latest checkpoint from: {latest_checkpoint_path}")
+            map_location = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            checkpoint = torch.load(latest_checkpoint_path, map_location=map_location)
+
+            # --- Config Compatibility Check (Optional but Recommended) ---
+            loaded_config = checkpoint.get("config", {})
+            # Compare key parameters that affect model architecture or data processing
+            critical_keys = ["dim", "n_layers", "n_heads", "ff_dim", "vocab_size", "max_seq_len", "tokenizer_name"]
+            mismatched_keys = []
+            if loaded_config:
+                for key in critical_keys:
+                    # Check if key exists in both and if they differ
+                    if key in loaded_config and key in CONFIG and loaded_config[key] != CONFIG[key]:
+                        mismatched_keys.append((key, loaded_config[key], CONFIG[key]))
+                    # Check if key missing in current config but present in checkpoint
+                    elif key in loaded_config and key not in CONFIG:
+                         mismatched_keys.append((key, loaded_config[key], "Not in current CONFIG"))
+                     # Check if key missing in checkpoint config but present in current
+                    elif key not in loaded_config and key in CONFIG:
+                         mismatched_keys.append((key, "Not in loaded CONFIG", CONFIG[key]))
+
+
+                if mismatched_keys:
+                    logging.warning("--- CONFIG MISMATCH DETECTED ---")
+                    logging.warning(f"Checkpoint '{os.path.basename(latest_checkpoint_path)}' was saved with different critical parameters:")
+                    for key, loaded_val, current_val in mismatched_keys:
+                        logging.warning(f"  - {key}: Checkpoint='{loaded_val}', Current='{current_val}'")
+                    # Decide whether to proceed: raise error, warn, or try anyway
+                    # For now, just warn strongly. Loading might fail or lead to issues.
+                    logging.warning("Proceeding with loading, but results may be unexpected or errors may occur.")
+            else:
+                logging.warning("Checkpoint does not contain configuration info. Cannot check compatibility.")
+            # --- End Config Check ---
+
+
+            try:
+                 # Strict=False can sometimes help load partially, but hides potential issues
+                 model.load_state_dict(checkpoint['model'], strict=True)
+            except RuntimeError as e:
+                 logging.error(f"Failed to load model state_dict: {e}")
+                 logging.error("This often happens due to architecture mismatch (check CONFIG) or corrupted checkpoint.")
+                 logging.error("Starting training from scratch.")
+                 return 0 # Cannot resume if model loading fails
+
+            try:
+                 optimizer.load_state_dict(checkpoint['optimizer'])
+            except ValueError as e:
+                 logging.warning(f"Could not load optimizer state_dict: {e}. Optimizer state will be reset.")
+                 # Reinitialize optimizer if state doesn't match? Or just proceed with current state.
+                 # Resetting optimizer state is safer if parameters changed.
+                 optimizer.state = defaultdict(dict) # Reset state
+                 logging.warning("Optimizer state reset.")
+            except Exception as e:
+                 logging.error(f"Unexpected error loading optimizer state: {e}. Starting training from scratch.")
+                 return 0
+
+            start_step = checkpoint.get('step', 0)
+            # Ensure step is non-negative, resume from next step
+            start_step = max(0, start_step) + 1 if isinstance(start_step, int) else 0
+
+
+            logging.info(f"Checkpoint loaded successfully. Resuming from optimizer step {start_step}.")
+            # Move optimizer state tensors to the correct device
+            for state in optimizer.state.values():
+                for k, v in state.items():
+                    if isinstance(v, torch.Tensor):
+                        try:
+                            state[k] = v.to(map_location)
+                        except Exception as e:
+                             logging.error(f"Failed to move optimizer tensor '{k}' to device '{map_location}': {e}")
+            return start_step
+
+        except FileNotFoundError:
+            logging.info(f"No checkpoint directory '{self.checkpoint_dir}' or files found. Starting training from scratch.")
+            return 0
+        except Exception as e:
+            logging.error(f"Error loading checkpoint from '{self.checkpoint_dir}': {e}. Starting training from scratch.")
+            # Clean up potentially partially loaded model/optimizer?
+            # Re-initializing might be safer depending on where the error occurred.
+            # For simplicity, we just return 0 here.
+            return 0
+
+
+# --- Training Function ---
+
+def train():
+    logging.info(f"Starting HROM training process on combined datasets (daily_dialog, empathetic_dialogues, blended_skill_talk, AlekseyKorshuk/persona-chat, elapt1c/Universal-Transformers-RawConvo, papahawk/conversational-01) using {CONFIG['num_gpus']} GPUs...") # Updated log message
+    logging.info(f"Configuration: {CONFIG}")
+
+    # --- Tokenizer Setup ---
+    tokenizer_trainer = TokenizerTrainer()
+    tokenizer_path = tokenizer_trainer.tokenizer_path
+    if not os.path.exists(tokenizer_path):
+        logging.info(f"Combined tokenizer '{CONFIG['tokenizer_name']}' not found. Training tokenizer...")
+        try:
+            # Pass trust_remote_code=True to load_dataset calls inside tokenizer training
+            tokenizer_trainer.train(CONFIG["datasets"])
+        except Exception as e:
+             logging.error(f"Failed during tokenizer training: {e}", exc_info=True)
+             return # Cannot proceed without a tokenizer
+    else:
+        logging.info(f"Loading existing combined tokenizer from {tokenizer_path}")
+    # Load the tokenizer instance *once* here for shared use
+    try:
+        tokenizer = tokenizer_trainer.get_tokenizer()
+        # Update CONFIG with actual token IDs (useful for downstream)
+        CONFIG['pad_token_id'] = tokenizer.token_to_id("<pad>")
+        CONFIG['bos_token_id'] = tokenizer.token_to_id("<s>")
+        CONFIG['eos_token_id'] = tokenizer.token_to_id("</s>")
+        logging.info(f"Loaded tokenizer. Vocab size: {tokenizer.get_vocab_size()}. Special IDs: PAD={CONFIG['pad_token_id']}, BOS={CONFIG['bos_token_id']}, EOS={CONFIG['eos_token_id']}")
+    except (FileNotFoundError, ValueError) as e:
+         logging.error(f"Failed to load tokenizer: {e}. Cannot continue.")
+         return
+
+    # --- Model Initialization ---
+    logging.info("Initializing HROM model...")
+    # Ensure vocab_size in config matches tokenizer
+    if CONFIG['vocab_size'] != tokenizer.get_vocab_size():
+         logging.warning(f"Config vocab_size ({CONFIG['vocab_size']}) differs from tokenizer vocab size ({tokenizer.get_vocab_size()}). Using tokenizer's size.")
+         CONFIG['vocab_size'] = tokenizer.get_vocab_size()
+    model = HROM()
+
+    # --- Calculate and Log Model Parameters ---
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    logging.info(f"Model initialized. Total parameters: {total_params:,}")
+    logging.info(f"Trainable parameters: {trainable_params:,}")
+    logging.info(f"Parameters (Millions): Total={total_params/1e6:.2f}M, Trainable={trainable_params/1e6:.2f}M")
+
+
+    # --- Dataset and DataLoader ---
+    logging.info("Setting up combined dataset and dataloader...")
+    try:
+         logging.info("Pre-loading/caching datasets...")
+         for ds_name in CONFIG["datasets"]:
+              logging.info(f"Checking cache for '{ds_name}'...")
+              try:
+                  # Load just the first example to trigger download/cache check
+                  _ = load_dataset(ds_name, split="train[:1]", download_mode="reuse_cache_if_exists", trust_remote_code=True) # Add trust_remote_code
+              except Exception as e:
+                  # Log error but try to continue, main dataset loading will handle final error
+                  logging.error(f"Could not pre-check dataset '{ds_name}': {e}")
+         logging.info("Dataset download/cache check presumed complete.")
+
+         # Pass the already loaded tokenizer instance
+         dataset = CombinedChatDataset(tokenizer)
+
+         # Check if dataset is empty after processing
+         if len(dataset) == 0:
+             logging.error("Dataset is empty after processing all sources. Cannot train.")
+             return
+
+         dataloader = DataLoader(
+             dataset,
+             batch_size=CONFIG["batch_size"],
+             collate_fn=CombinedChatDataset.collate_fn, # Use static method
+             shuffle=True,
+             # Adjust num_workers based on available cores, be conservative
+             num_workers=min(4 * CONFIG["num_gpus"], os.cpu_count() // 2 if (os.cpu_count() and os.cpu_count() > 1) else 1), # Increased num_workers
+             pin_memory=torch.cuda.is_available(),
+             prefetch_factor=2 if torch.cuda.is_available() and os.cpu_count() and os.cpu_count() > 1 else None,
+             drop_last=False # Keep last batch even if smaller
+             )
+    except Exception as e:
+         logging.error(f"Failed to initialize dataset/dataloader: {e}", exc_info=True)
+         return
+
+    # --- Trainer, Checkpoint, Safety ---
+    logging.info("Initializing Trainer, Checkpoint Manager, and Safety Manager...")
+    # Pass the loaded tokenizer instance and main_device to SafetyManager
+    trainer_obj = HROMTrainer(model, tokenizer)
+    checkpoint_manager = CheckpointManager() # Uses CONFIG["checkpoint_dir"]
+    safety = SafetyManager(trainer_obj.model, tokenizer, trainer_obj.main_device) # Pass main_device to SafetyManager
+
+    # --- Load Checkpoint ---
+    start_optimizer_step = checkpoint_manager.load_latest(trainer_obj.model, trainer_obj.optimizer) # Pass potentially DP wrapped model
+    # Ensure model is on correct device after loading (already handled in trainer init)
+
+    # --- Training Loop ---
+    logging.info(f"Starting training from optimizer step {start_optimizer_step}")
+    optimizer_step = start_optimizer_step
+    total_loss_accum = 0.0
+    # Calculate starting batch step based on loaded optimizer step and grad accum
+    batch_step = optimizer_step * CONFIG["grad_accum_steps"]
+    epochs_completed = batch_step // len(dataloader) if len(dataloader) > 0 else 0
+    start_epoch = epochs_completed # Start from the epoch corresponding to the loaded step
+
+    # Estimate total steps (can be useful for LR scheduling if implementing decay)
+    try:
+        if len(dataloader) == 0:
+            raise ValueError("DataLoader has zero length. Cannot estimate total steps.")
+        total_optimizer_steps = (len(dataloader) * CONFIG["num_epochs"]) // CONFIG["grad_accum_steps"]
+        logging.info(f"Estimated dataset size: {len(dataset)}")
+        logging.info(f"Estimated batches per epoch: {len(dataloader)}")
+        logging.info(f"Gradient Accumulation Steps: {CONFIG['grad_accum_steps']}")
+        logging.info(f"Effective Batch Size: {CONFIG['batch_size'] * CONFIG['grad_accum_steps']}")
+        logging.info(f"Target Epochs: {CONFIG['num_epochs']}")
+        logging.info(f"Estimated total optimizer steps for {CONFIG['num_epochs']} epochs: {total_optimizer_steps}")
+    except Exception as e:
+        logging.warning(f"Could not accurately estimate dataloader length or total steps: {e}")
+        total_optimizer_steps = -1 # Indicate unknown total steps
+
+
+    trainer_obj.model.train() # Ensure model is in training mode
+
+    for epoch in range(start_epoch, CONFIG["num_epochs"]):
+        logging.info(f"--- Starting Epoch {epoch+1}/{CONFIG['num_epochs']} ---")
+        epoch_loss = 0.0
+        num_batches_in_epoch = 0
+
+        # Use enumerate starting from 1 for batch count if preferred
+        for i, batch in enumerate(dataloader):
+            # Check if batch is valid (collate_fn might return None)
+            if batch is None:
+                 logging.warning(f"Skipping empty batch at step {i} in epoch {epoch+1}")
+                 continue
+
+            # Forward and backward pass (scaled loss)
+            loss = trainer_obj.train_step(batch)
+            if loss is None or torch.isnan(torch.tensor(loss)) or torch.isinf(torch.tensor(loss)):
+                 logging.error(f"NaN, Inf, or None loss detected: {loss}. Epoch {epoch+1}, Batch {i}, Opt Step {optimizer_step}. Stopping.")
+                 # Try saving a 'nan_inf' checkpoint before exiting
+                 checkpoint_manager.save(trainer_obj.model, trainer_obj.optimizer, f"{optimizer_step}_error") # Pass potentially DP wrapped model
+                 return
+
+            total_loss_accum += loss
+            epoch_loss += loss
+            num_batches_in_epoch += 1
+            batch_step += 1 # Increment global batch counter (tracks batches processed)
+
+            # Gradient Accumulation Check & Optimizer Step
+            # Check if it's time to perform an optimizer step
+            if batch_step % CONFIG["grad_accum_steps"] == 0:
+                current_lr = trainer_obj.clip_and_step(optimizer_step) # Pass current opt step for LR schedule
+
+                # Calculate average loss over accumulation steps for logging
+                avg_loss = total_loss_accum / CONFIG["grad_accum_steps"]
+                total_loss_accum = 0.0 # Reset loss accumulator
+
+                # Logging (use main_device for logging info)
+                if optimizer_step % CONFIG["debug_interval"] == 0:
+                    logging.info(f"Epoch {epoch+1} | Opt Step {optimizer_step} | Batch Step {batch_step} | Avg Loss: {avg_loss:.4f} | LR: {current_lr:.2e} | Device: {trainer_obj.main_device}") # Added device info to log
+                    # Trigger debug generation less frequently or based on condition
+                    if optimizer_step % (CONFIG["debug_interval"] * 5) == 0: # e.g., every 5 debug intervals
+                         safety.debug_generation("<user> Hi there! How are you doing today?") # Use a generic debug prompt
+
+                # Checkpointing
+                if optimizer_step > 0 and optimizer_step % CONFIG["checkpoint_interval"] == 0:
+                    logging.info(f"Checkpoint interval reached at optimizer step {optimizer_step}.")
+                    checkpoint_manager.save(trainer_obj.model, trainer_obj.optimizer, optimizer_step) # Pass potentially DP wrapped model
+                    # Optional: Run a generation check after saving checkpoint
+                    safety.debug_generation("<user> Hi! How are you?")
+
+                optimizer_step += 1 # Increment optimizer step count *after* performing the step
+
+        # --- End of Epoch ---
+        avg_epoch_loss = epoch_loss / num_batches_in_epoch if num_batches_in_epoch > 0 else 0
+        logging.info(f"--- Finished Epoch {epoch+1}/{CONFIG['num_epochs']} | Average Epoch Loss: {avg_epoch_loss:.4f} ---")
+
+        # Save checkpoint at the end of each epoch
+        checkpoint_manager.save(trainer_obj.model, trainer_obj.optimizer, f"epoch{epoch+1}_step{optimizer_step}") # Pass potentially DP wrapped model
+        # Optionally run debug generation at end of epoch
+        safety.debug_generation("<user> Hi! Whats up?")
+
+
+    logging.info(f"Training finished after {CONFIG['num_epochs']} target epochs.")
+    # Final save
+    logging.info("Saving final model state...")
+    checkpoint_manager.save(trainer_obj.model, trainer_obj.optimizer, f"final_step{optimizer_step}") # Pass potentially DP wrapped model
+
+
+if __name__ == "__main__":
+    # Ensures imports happen after setting the env var if script is run directly
+    train()