#!/usr/bin/env python3 """ Agent Token Management System This module provides comprehensive agent token management for multi-agent training, including special token handling, embedding management, and integration with existing tokenization systems. """ import os import json import logging from typing import Dict, List, Optional, Tuple, Any, Union from dataclasses import dataclass from pathlib import Path import torch from transformers import AutoTokenizer, PreTrainedTokenizer from transformers.tokenization_utils_base import PreTrainedTokenizerBase logger = logging.getLogger(__name__) @dataclass class AgentTokenConfig: """Configuration for agent token management""" agent_prefix: str = "<|agent:" agent_suffix: str = "|>" special_tokens: Optional[Dict[str, str]] = None resize_embeddings: bool = True save_tokens: bool = True tokens_file: str = "agent_tokens.json" class AgentTokenManager: """ Manages agent-specific tokens and their integration with tokenizers """ def __init__(self, config: AgentTokenConfig): self.config = config self.agent_tokens: Dict[str, str] = {} self.token_ids: Dict[str, int] = {} self.original_vocab_size: Optional[int] = None def generate_agent_tokens(self, agents: List[str]) -> List[str]: """Generate agent tokens for given agent list""" tokens = [] for agent in agents: token = f"{self.config.agent_prefix}{agent}{self.config.agent_suffix}" tokens.append(token) self.agent_tokens[agent] = token logger.info(f"Generated {len(tokens)} agent tokens: {tokens}") return tokens def add_agent_tokens_to_tokenizer(self, tokenizer: PreTrainedTokenizer, agents: List[str]) -> Tuple[PreTrainedTokenizer, List[str]]: """ Add agent tokens to tokenizer and return updated tokenizer with token list """ if not agents: logger.warning("No agents provided, skipping token addition") return tokenizer, [] # Generate agent tokens agent_tokens = self.generate_agent_tokens(agents) # Check which tokens need to be added existing_tokens = set(tokenizer.get_vocab().keys()) tokens_to_add = [token for token in agent_tokens if token not in existing_tokens] if not tokens_to_add: logger.info("All agent tokens already exist in tokenizer") return tokenizer, agent_tokens # Store original vocab size self.original_vocab_size = len(tokenizer) # Add special tokens logger.info(f"Adding {len(tokens_to_add)} new agent tokens to tokenizer") tokenizer.add_special_tokens({ "additional_special_tokens": tokens_to_add }) # Update token IDs mapping for agent, token in self.agent_tokens.items(): if token in tokenizer.get_vocab(): self.token_ids[agent] = tokenizer.convert_tokens_to_ids(token) logger.info(f"Added agent tokens. New vocab size: {len(tokenizer)}") return tokenizer, agent_tokens def resize_model_embeddings(self, model: torch.nn.Module, tokenizer: PreTrainedTokenizer) -> torch.nn.Module: """ Resize model embeddings to accommodate new agent tokens """ if not self.config.resize_embeddings: logger.info("Embedding resize disabled, skipping") return model if self.original_vocab_size is None: logger.warning("Original vocab size not set, cannot resize embeddings") return model new_vocab_size = len(tokenizer) if new_vocab_size == self.original_vocab_size: logger.info("Vocab size unchanged, no embedding resize needed") return model logger.info(f"Resizing model embeddings from {self.original_vocab_size} to {new_vocab_size}") # Resize embeddings model.resize_token_embeddings(new_vocab_size) # Initialize new embeddings (copy from unk token or use random initialization) if hasattr(model, 'get_input_embeddings'): embeddings = model.get_input_embeddings() if hasattr(embeddings, 'weight'): with torch.no_grad(): # Initialize new embeddings with small random values new_embeddings = embeddings.weight[self.original_vocab_size:] torch.nn.init.normal_(new_embeddings, mean=0.0, std=0.02) logger.info("Model embeddings resized successfully") return model def format_agent_prompt(self, agent: str, text: str) -> str: """Format text with agent token prefix""" if agent not in self.agent_tokens: logger.warning(f"Agent '{agent}' not found in token mappings") return text agent_token = self.agent_tokens[agent] return f"{agent_token}\n{text}" def extract_agent_from_text(self, text: str) -> Optional[str]: """Extract agent name from text if it starts with agent token""" for agent, token in self.agent_tokens.items(): if text.startswith(token): return agent return None def get_agent_token_id(self, agent: str) -> Optional[int]: """Get token ID for agent token""" return self.token_ids.get(agent) def save_agent_tokens(self, output_dir: str) -> str: """Save agent tokens to file""" if not self.config.save_tokens: return "" os.makedirs(output_dir, exist_ok=True) tokens_file = os.path.join(output_dir, self.config.tokens_file) tokens_data = { "agent_tokens": self.agent_tokens, "token_ids": self.token_ids, "config": { "agent_prefix": self.config.agent_prefix, "agent_suffix": self.config.agent_suffix, "original_vocab_size": self.original_vocab_size } } with open(tokens_file, 'w') as f: json.dump(tokens_data, f, indent=2) logger.info(f"Saved agent tokens to {tokens_file}") return tokens_file def load_agent_tokens(self, tokens_file: str) -> bool: """Load agent tokens from file""" if not os.path.isfile(tokens_file): logger.warning(f"Agent tokens file not found: {tokens_file}") return False try: with open(tokens_file, 'r') as f: tokens_data = json.load(f) self.agent_tokens = tokens_data.get("agent_tokens", {}) self.token_ids = tokens_data.get("token_ids", {}) config_data = tokens_data.get("config", {}) self.original_vocab_size = config_data.get("original_vocab_size") logger.info(f"Loaded {len(self.agent_tokens)} agent tokens from {tokens_file}") return True except Exception as e: logger.error(f"Failed to load agent tokens: {e}") return False def get_agent_statistics(self) -> Dict[str, Any]: """Get statistics about agent tokens""" return { "total_agents": len(self.agent_tokens), "agents": list(self.agent_tokens.keys()), "token_ids": self.token_ids, "original_vocab_size": self.original_vocab_size, "config": { "agent_prefix": self.config.agent_prefix, "agent_suffix": self.config.agent_suffix } } class AgentTokenizer: """ Enhanced tokenizer wrapper that integrates agent token management """ def __init__(self, tokenizer: PreTrainedTokenizer, agent_manager: AgentTokenManager): self.tokenizer = tokenizer self.agent_manager = agent_manager def tokenize_agent_text(self, agent: str, text: str, **kwargs) -> Dict[str, Any]: """Tokenize text with agent prefix""" formatted_text = self.agent_manager.format_agent_prompt(agent, text) return self.tokenizer(formatted_text, **kwargs) def decode_agent_tokens(self, token_ids: Union[List[int], torch.Tensor], **kwargs) -> str: """Decode token IDs back to text""" return self.tokenizer.decode(token_ids, **kwargs) def get_agent_attention_mask(self, input_ids: torch.Tensor, agent: str) -> torch.Tensor: """Get attention mask with special handling for agent tokens""" attention_mask = torch.ones_like(input_ids) # Find agent token position agent_token_id = self.agent_manager.get_agent_token_id(agent) if agent_token_id is not None: # Ensure agent token is attended to agent_positions = (input_ids == agent_token_id) attention_mask[agent_positions] = 1 return attention_mask def __getattr__(self, name): """Delegate unknown attributes to underlying tokenizer""" return getattr(self.tokenizer, name) class AgentTokenValidator: """Validator for agent token configurations""" @staticmethod def validate_agent_tokens(agents: List[str], config: AgentTokenConfig) -> Dict[str, Any]: """Validate agent token configuration""" validation_result = { "valid": True, "errors": [], "warnings": [], "tokens": {} } if not agents: validation_result["warnings"].append("No agents provided") return validation_result # Check for duplicate agents if len(agents) != len(set(agents)): validation_result["errors"].append("Duplicate agents found") validation_result["valid"] = False # Generate and validate tokens manager = AgentTokenManager(config) tokens = manager.generate_agent_tokens(agents) # Check for token conflicts token_set = set(tokens) if len(token_set) != len(tokens): validation_result["errors"].append("Duplicate tokens generated") validation_result["valid"] = False # Check token length for agent, token in zip(agents, tokens): if len(token) > 50: # Reasonable limit validation_result["warnings"].append(f"Long token for agent '{agent}': {token}") validation_result["tokens"] = dict(zip(agents, tokens)) return validation_result @staticmethod def validate_tokenizer_compatibility(tokenizer: PreTrainedTokenizer, agents: List[str], config: AgentTokenConfig) -> Dict[str, Any]: """Validate tokenizer compatibility with agent tokens""" validation_result = { "compatible": True, "errors": [], "warnings": [], "existing_tokens": [], "new_tokens": [] } if not agents: return validation_result # Generate tokens manager = AgentTokenManager(config) tokens = manager.generate_agent_tokens(agents) # Check existing vocabulary vocab = tokenizer.get_vocab() for agent, token in zip(agents, tokens): if token in vocab: validation_result["existing_tokens"].append(agent) else: validation_result["new_tokens"].append(agent) # Check for potential conflicts for token in tokens: if token in vocab: # Check if it's already a special token if hasattr(tokenizer, 'special_tokens_map'): special_tokens = tokenizer.special_tokens_map if token not in special_tokens.values(): validation_result["warnings"].append(f"Token '{token}' exists in vocab but not as special token") return validation_result # Integration with existing MoE framework class MoEAgentTokenIntegration: """ Integration layer between agent tokens and MoE framework """ def __init__(self, agent_manager: AgentTokenManager): self.agent_manager = agent_manager self.agent_to_expert_mapping: Dict[str, str] = {} def map_agent_to_expert(self, agent: str, expert: str): """Map agent to MoE expert specialization""" self.agent_to_expert_mapping[agent] = expert logger.info(f"Mapped agent '{agent}' to expert '{expert}'") def get_expert_for_agent(self, agent: str) -> Optional[str]: """Get expert specialization for agent""" return self.agent_to_expert_mapping.get(agent) def format_moe_prompt(self, agent: str, text: str, expert: Optional[str] = None) -> str: """Format prompt for MoE framework with agent and expert context""" # Start with agent token formatted_text = self.agent_manager.format_agent_prompt(agent, text) # Add expert context if available if expert: expert_context = f"\n<|expert:{expert}|>\n" formatted_text = formatted_text.replace("\n", expert_context, 1) return formatted_text def extract_agent_and_expert(self, text: str) -> Tuple[Optional[str], Optional[str]]: """Extract both agent and expert from formatted text""" agent = self.agent_manager.extract_agent_from_text(text) # Extract expert if present expert = None if "<|expert:" in text and "|>" in text: start = text.find("<|expert:") + 9 end = text.find("|>", start) if end > start: expert = text[start:end] return agent, expert # Example usage and testing if __name__ == "__main__": # Configure logging logging.basicConfig(level=logging.INFO) # Example configuration config = AgentTokenConfig( agent_prefix="<|agent:", agent_suffix="|>", resize_embeddings=True ) # Example agents agents = ["SWE", "SQE", "DevOps", "Architect", "Security"] # Create agent manager manager = AgentTokenManager(config) # Generate tokens tokens = manager.generate_agent_tokens(agents) print(f"Generated tokens: {tokens}") # Example tokenizer (would be loaded from actual model) from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-MoE-instruct") # Add tokens to tokenizer updated_tokenizer, agent_tokens = manager.add_agent_tokens_to_tokenizer(tokenizer, agents) print(f"Updated tokenizer vocab size: {len(updated_tokenizer)}") print(f"Agent token IDs: {manager.token_ids}") # Test formatting test_text = "How do I implement a binary search?" formatted = manager.format_agent_prompt("SWE", test_text) print(f"Formatted prompt: {formatted}") # Test extraction extracted_agent = manager.extract_agent_from_text(formatted) print(f"Extracted agent: {extracted_agent}")