""" Memory Training Module for MangoMAS Local This module implements specialized training for memory and context retention capabilities, adapted from the AWS backup system for local training. """ import json import logging import os import random from typing import Any, Dict, List import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset from ..core_framework import SpecializedTrainingModule, TrainingModuleConfig logger = logging.getLogger(__name__) class MemoryDataset(Dataset): """Dataset for training memory and context retention capabilities.""" def __init__(self, data_path: str, tokenizer, max_length: int = 1024): """ Initialize the memory dataset. Args: data_path: Path to the memory training data file tokenizer: Tokenizer for text processing max_length: Maximum sequence length """ self.tokenizer = tokenizer self.max_length = max_length self.data = self._load_data(data_path) logger.info(f"Loaded memory dataset with {len(self.data)} examples") def _load_data(self, data_path: str) -> List[Dict]: """Load memory training data.""" data = [] with open(data_path, "r", encoding="utf-8") as f: for line in f: try: item = json.loads(line.strip()) # Validate required fields if "conversation" in item and isinstance( item["conversation"], list ): data.append(item) except json.JSONDecodeError: continue return data def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] # Format the conversation for memory training conversation = item["conversation"] context = "\n".join( [f"{turn['role']}: {turn['content']}" for turn in conversation[:-1]] ) target = conversation[-1]["content"] prompt = f"Context:\n{context}\nResponse: {target}" # Tokenize encoding = self.tokenizer( prompt, max_length=self.max_length, padding="max_length", truncation=True, return_tensors="pt", ) return { "input_ids": encoding["input_ids"].squeeze(), "attention_mask": encoding["attention_mask"].squeeze(), "labels": encoding["input_ids"].squeeze(), } class MemoryTrainingModule(SpecializedTrainingModule): """Specialized training module for memory and context retention capabilities.""" def __init__(self, config: TrainingModuleConfig, tokenizer): """ Initialize the memory training module. Args: config: Module configuration tokenizer: Tokenizer for text processing """ super().__init__(config, tokenizer) # Initialize memory-specific components self.memory_loss = nn.CrossEntropyLoss(ignore_index=-100) self.metrics = { "memory_loss": 0.0, "context_retention": 0.0, "coherence_score": 0.0, } logger.info("Initialized MemoryTrainingModule") def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: """ Prepare a batch of data for memory training. Args: batch: The input batch from the dataloader Returns: Processed batch ready for memory training """ # Move batch to device prepared_batch = {} for key, value in batch.items(): if isinstance(value, torch.Tensor): prepared_batch[key] = value.to(self.device) else: prepared_batch[key] = value return prepared_batch def compute_loss( self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor] ) -> torch.Tensor: """ Compute the memory-specific loss. Args: student_outputs: Outputs from the student model teacher_outputs: Outputs from the teacher model batch: The processed input batch Returns: Loss tensor for memory training """ try: # Extract logits from model outputs if hasattr(student_outputs, "logits"): student_logits = student_outputs.logits else: student_logits = student_outputs if hasattr(teacher_outputs, "logits"): teacher_logits = teacher_outputs.logits else: teacher_logits = teacher_outputs # Get labels from batch labels = batch.get("labels", batch.get("input_ids")) # Compute cross entropy loss for memory shift_logits = student_logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() memory_loss = self.memory_loss( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) ) # Add KL divergence loss between student and teacher if teacher_logits is not None: kl_loss = F.kl_div( F.log_softmax(student_logits, dim=-1), F.softmax(teacher_logits, dim=-1), reduction="batchmean", ) total_loss = memory_loss + 0.1 * kl_loss else: total_loss = memory_loss # Update metrics self.metrics["memory_loss"] = memory_loss.item() return total_loss * self.loss_weight except Exception as e: logger.error(f"Error computing memory loss: {e}") # Return a small loss to avoid training failure return torch.tensor(0.01, requires_grad=True) def get_metrics(self) -> Dict[str, float]: """ Get metrics specific to memory training. Returns: Dictionary of memory metrics """ return self.metrics.copy() def generate_synthetic_memory_data( self, output_path: str, num_samples: int = 1000 ) -> None: """ Generate synthetic memory training data. Args: output_path: Path to save the generated data num_samples: Number of samples to generate """ # This is a simplified implementation based on the AWS backup # In a full implementation, this would be much more sophisticated conversation_templates = [ [ { "role": "user", "content": "Hi, my name is Alex and I'm interested in machine learning.", }, { "role": "assistant", "content": "Hello Alex! I'd be happy to discuss machine learning with you. What aspects are you most interested in?", }, { "role": "user", "content": "I'm particularly interested in natural language processing.", }, { "role": "assistant", "content": "NLP is a fascinating field! It's used for tasks like translation, summarization, and question answering.", }, { "role": "user", "content": "What do you think would be a good first project?", }, { "role": "assistant", "content": "For a beginner in NLP, I'd recommend starting with a text classification project, like sentiment analysis.", }, ], [ { "role": "user", "content": "I'm planning a trip to Japan next spring.", }, { "role": "assistant", "content": "That sounds exciting! Japan is beautiful in spring with cherry blossoms. What cities are you planning to visit?", }, { "role": "user", "content": "I'm thinking Tokyo, Kyoto, and maybe Osaka.", }, { "role": "assistant", "content": "Great choices! Tokyo has modern attractions, Kyoto has historical temples, and Osaka is known for amazing food.", }, { "role": "user", "content": "What's the best way to travel between these cities?", }, { "role": "assistant", "content": "The Shinkansen (bullet train) is the most efficient way to travel between these cities. It's fast, comfortable, and reliable.", }, ], ] recall_templates = [ { "recall_context": "what was my name again?", "recall_target": "Your name is Alex, as you mentioned at the beginning of our conversation.", }, { "recall_context": "which cities did I say I wanted to visit?", "recall_target": "You mentioned you're planning to visit Tokyo, Kyoto, and possibly Osaka during your trip to Japan.", }, ] # Generate variations output_data = [] for _ in range(num_samples): template_idx = random.randint(0, len(conversation_templates) - 1) conversation = conversation_templates[template_idx].copy() # Add a recall question if this is the right template if template_idx < len(recall_templates): recall_template = recall_templates[template_idx] # Add a user question asking for recall conversation.append( {"role": "user", "content": recall_template["recall_context"]} ) # Create the full example with recall targets example = { "conversation": conversation, "recall_context": recall_template["recall_context"], "recall_target": recall_template["recall_target"], "metadata": {"generated": True, "requires_memory": True}, } else: # Regular conversation without specific recall target example = { "conversation": conversation, "metadata": {"generated": True, "requires_memory": False}, } output_data.append(example) # Save to file os.makedirs(os.path.dirname(output_path), exist_ok=True) with open(output_path, "w", encoding="utf-8") as f: for item in output_data: f.write(json.dumps(item) + "\n") logger.info( f"Generated {len(output_data)} synthetic memory examples at {output_path}" )