Upload mutators.py

mutators.py

@@ -0,0 +1,163 @@
+"""
+Mutation operators for biological language models.
+Adds controlled noise and structural perturbations to model weights.
+"""
+
+import os
+import torch
+import numpy as np
+import logging
+from typing import Optional
+from transformers import AutoModelForMaskedLM
+from safetensors.torch import save_file, load_file
+
+logger = logging.getLogger(__name__)
+
+
+class BioMutator:
+    """Base mutation operator."""
+    
+    def mutate(self, model_path: str, save_path: str) -> str:
+        """Apply mutation to a model and save result."""
+        raise NotImplementedError
+
+
+class GaussianNoiseMutator(BioMutator):
+    """Add Gaussian noise to model weights."""
+    
+    def __init__(self, std: float = 0.01, layer_filter: Optional[str] = None):
+        self.std = std
+        self.layer_filter = layer_filter  # If set, only mutate matching layers
+    
+    def mutate(self, model_path: str, save_path: str) -> str:
+        logger.info(f"Applying Gaussian noise mutation (std={self.std}) to {model_path}")
+        
+        try:
+            state = load_file(os.path.join(model_path, "model.safetensors"))
+        except:
+            state = torch.load(
+                os.path.join(model_path, "pytorch_model.bin"),
+                map_location="cpu",
+            )
+        
+        mutated = {}
+        for key, tensor in state.items():
+            if self.layer_filter is not None and self.layer_filter not in key:
+                mutated[key] = tensor.clone()
+                continue
+            
+            # Generate noise matching tensor shape
+            noise = torch.randn_like(tensor) * self.std
+            mutated[key] = tensor + noise
+        
+        # Save
+        os.makedirs(save_path, exist_ok=True)
+        save_file(mutated, os.path.join(save_path, "model.safetensors"))
+        
+        # Copy auxiliary files
+        self._copy_auxiliary_files(model_path, save_path)
+        
+        logger.info(f"Mutated model saved to {save_path}")
+        return save_path
+    
+    def _copy_auxiliary_files(self, source_path: str, dest_path: str) -> None:
+        import shutil
+        
+        for fname in ["config.json", "tokenizer.json", "tokenizer_config.json", "vocab.txt"]:
+            src = os.path.join(source_path, fname)
+            if os.path.exists(src):
+                shutil.copy2(src, os.path.join(dest_path, fname))
+
+
+class LayerScaleMutator(BioMutator):
+    """Scale specific layers by a random factor."""
+    
+    def __init__(self, scale_range: tuple = (0.9, 1.1), target_layers: Optional[list] = None):
+        self.scale_range = scale_range
+        self.target_layers = target_layers  # e.g., ["encoder.layer.0", "encoder.layer.1"]
+    
+    def mutate(self, model_path: str, save_path: str) -> str:
+        logger.info(f"Applying layer scale mutation to {model_path}")
+        
+        try:
+            state = load_file(os.path.join(model_path, "model.safetensors"))
+        except:
+            state = torch.load(
+                os.path.join(model_path, "pytorch_model.bin"),
+                map_location="cpu",
+            )
+        
+        mutated = {}
+        for key, tensor in state.items():
+            # Check if this key matches a target layer
+            should_scale = False
+            if self.target_layers:
+                for layer_name in self.target_layers:
+                    if layer_name in key:
+                        should_scale = True
+                        break
+            else:
+                # Randomly select ~10% of layers to scale
+                should_scale = np.random.random() < 0.1
+            
+            if should_scale:
+                scale = np.random.uniform(*self.scale_range)
+                mutated[key] = tensor * scale
+                logger.debug(f"Scaled {key} by {scale:.4f}")
+            else:
+                mutated[key] = tensor.clone()
+        
+        os.makedirs(save_path, exist_ok=True)
+        save_file(mutated, os.path.join(save_path, "model.safetensors"))
+        
+        self._copy_auxiliary_files(model_path, save_path)
+        
+        logger.info(f"Layer-scaled model saved to {save_path}")
+        return save_path
+    
+    def _copy_auxiliary_files(self, source_path: str, dest_path: str) -> None:
+        import shutil
+        
+        for fname in ["config.json", "tokenizer.json", "tokenizer_config.json", "vocab.txt"]:
+            src = os.path.join(source_path, fname)
+            if os.path.exists(src):
+                shutil.copy2(src, os.path.join(dest_path, fname))
+
+
+class DropoutMutator(BioMutator):
+    """Randomly zero out a fraction of weights (structural pruning)."""
+    
+    def __init__(self, dropout_rate: float = 0.01):
+        self.dropout_rate = dropout_rate
+    
+    def mutate(self, model_path: str, save_path: str) -> str:
+        logger.info(f"Applying dropout mutation (rate={self.dropout_rate}) to {model_path}")
+        
+        try:
+            state = load_file(os.path.join(model_path, "model.safetensors"))
+        except:
+            state = torch.load(
+                os.path.join(model_path, "pytorch_model.bin"),
+                map_location="cpu",
+            )
+        
+        mutated = {}
+        for key, tensor in state.items():
+            mask = torch.rand_like(tensor) > self.dropout_rate
+            mutated[key] = tensor * mask
+        
+        os.makedirs(save_path, exist_ok=True)
+        save_file(mutated, os.path.join(save_path, "model.safetensors"))
+        
+        self._copy_auxiliary_files(model_path, save_path)
+        
+        logger.info(f"Dropout-mutated model saved to {save_path}")
+        return save_path
+    
+    def _copy_auxiliary_files(self, source_path: str, dest_path: str) -> None:
+        import shutil
+        
+        for fname in ["config.json", "tokenizer.json", "tokenizer_config.json", "vocab.txt"]:
+            src = os.path.join(source_path, fname)
+            if os.path.exists(src):
+                shutil.copy2(src, os.path.join(dest_path, fname))