Upload evaluator.py

evaluator.py

@@ -0,0 +1,409 @@
+"""
+Evaluator for biological language models on synthetic sequence tasks.
+Supports masked language models (ESM-2, NT) and autoregressive models.
+"""
+
+import re
+import logging
+from typing import List, Dict, Optional
+import numpy as np
+from transformers import AutoModelForMaskedLM, AutoTokenizer, EsmTokenizer
+import torch
+from difflib import SequenceMatcher
+
+from .tasks import BioTask
+
+logger = logging.getLogger(__name__)
+
+
+class BioEvaluator:
+    """Evaluates biological language models on sequence tasks."""
+    
+    def __init__(self, device: str = "auto", max_length: int = 1024):
+        self.device = device if device != "auto" else ("cuda" if torch.cuda.is_available() else "cpu")
+        self.max_length = max_length
+        self._model_cache = {}
+        self._tokenizer_cache = {}
+    
+    def _load_model(self, model_path: str):
+        """Load model with caching."""
+        if model_path not in self._model_cache:
+            logger.info(f"Loading model from {model_path}")
+            try:
+                model = AutoModelForMaskedLM.from_pretrained(
+                    model_path,
+                    torch_dtype=torch.bfloat16,
+                    trust_remote_code=True,
+                )
+            except:
+                # Fallback if not standard masked LM
+                from transformers import AutoModel
+                model = AutoModel.from_pretrained(
+                    model_path,
+                    torch_dtype=torch.bfloat16,
+                    trust_remote_code=True,
+                )
+            
+            model = model.to(self.device)
+            model.eval()
+            self._model_cache[model_path] = model
+        
+        return self._model_cache[model_path]
+    
+    def _load_tokenizer(self, model_path: str):
+        """Load tokenizer with caching."""
+        if model_path not in self._tokenizer_cache:
+            logger.info(f"Loading tokenizer from {model_path}")
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_path,
+                trust_remote_code=True,
+            )
+            self._tokenizer_cache[model_path] = tokenizer
+        
+        return self._tokenizer_cache[model_path]
+    
+    def evaluate_model(
+        self,
+        model_path: str,
+        tasks: List[BioTask],
+    ) -> Dict[str, float]:
+        """Evaluate a model on a list of tasks. Returns task_id -> score mapping."""
+        model = self._load_model(model_path)
+        tokenizer = self._load_tokenizer(model_path)
+        
+        results = {}
+        
+        for task in tasks:
+            try:
+                score = self._evaluate_single_task(model, tokenizer, task)
+                results[task.task_id] = score
+            except Exception as e:
+                logger.error(f"Error evaluating task {task.task_id}: {e}")
+                results[task.task_id] = 0.0
+        
+        return results
+    
+    def _evaluate_single_task(
+        self,
+        model: torch.nn.Module,
+        tokenizer,
+        task: BioTask,
+    ) -> float:
+        """Evaluate a single task."""
+        
+        if task.evaluation_metric == "sequence_identity":
+            return self._eval_sequence_identity(model, tokenizer, task)
+        
+        elif task.evaluation_metric == "sequence_similarity":
+            return self._eval_sequence_similarity(model, tokenizer, task)
+        
+        elif task.evaluation_metric == "contains_substring":
+            return self._eval_contains_substring(model, tokenizer, task)
+        
+        elif task.evaluation_metric == "exact_match":
+            return self._eval_exact_match(model, tokenizer, task)
+        
+        elif task.evaluation_metric == "perplexity":
+            return self._eval_perplexity(model, tokenizer, task)
+        
+        elif task.evaluation_metric == "rna_structure_similarity":
+            return self._eval_rna_structure(model, tokenizer, task)
+        
+        else:
+            logger.warning(f"Unknown metric: {task.evaluation_metric}, defaulting to sequence similarity")
+            return self._eval_sequence_similarity(model, tokenizer, task)
+    
+    def _get_model_output(self, model, tokenizer, prompt: str) -> str:
+        """Get model output for a prompt."""
+        # For masked LMs, we use the masked prediction approach
+        # For autoregressive models, we'd use generation
+        
+        if task_has_mask := "<mask>" in prompt or "[MASK]" in prompt:
+            # Masked prediction task
+            return self._predict_masked(model, tokenizer, prompt)
+        else:
+            # For sequence continuation, try autoregressive generation if model supports it
+            return self._generate_sequence(model, tokenizer, prompt)
+    
+    def _predict_masked(self, model, tokenizer, prompt: str) -> str:
+        """Predict masked tokens in a sequence."""
+        # Tokenize
+        tokens = tokenizer.tokenize(prompt)
+        
+        # Find mask positions
+        mask_token = tokenizer.mask_token or "<mask>"
+        mask_positions = [i for i, t in enumerate(tokens) if t == mask_token or t == "[MASK]"]
+        
+        if not mask_positions:
+            # No mask found, just return prompt
+            return prompt
+        
+        # Convert to IDs
+        input_ids = tokenizer.encode(prompt, return_tensors="pt", max_length=self.max_length, truncation=True)
+        input_ids = input_ids.to(self.device)
+        
+        # Get predictions
+        with torch.no_grad():
+            outputs = model(input_ids)
+            logits = outputs.logits
+        
+        # Fill in masks
+        predicted_tokens = tokens.copy()
+        for pos in mask_positions:
+            mask_logits = logits[0, pos + 1]  # +1 for CLS if present
+            predicted_id = torch.argmax(mask_logits).item()
+            predicted_token = tokenizer.convert_ids_to_tokens([predicted_id])[0]
+            predicted_tokens[pos] = predicted_token
+        
+        # Reconstruct
+        return tokenizer.convert_tokens_to_string(predicted_tokens)
+    
+    def _generate_sequence(self, model, tokenizer, prompt: str, max_new_tokens: int = 50) -> str:
+        """Generate a sequence continuation."""
+        # Simple greedy generation for masked LM models
+        # For true autoregressive models, this would use generate()
+        
+        input_ids = tokenizer.encode(prompt, return_tensors="pt", max_length=self.max_length, truncation=True)
+        input_ids = input_ids.to(self.device)
+        
+        generated = input_ids.clone()
+        
+        # Greedy token-by-token generation
+        for _ in range(max_new_tokens):
+            with torch.no_grad():
+                outputs = model(generated)
+                logits = outputs.logits
+            
+            # Get next token prediction
+            next_token_logits = logits[0, -1, :]
+            next_token_id = torch.argmax(next_token_logits).item()
+            
+            # Append
+            next_token = torch.tensor([[next_token_id]], device=self.device)
+            generated = torch.cat([generated, next_token], dim=1)
+            
+            # Check for EOS
+            if next_token_id == tokenizer.eos_token_id:
+                break
+        
+        return tokenizer.decode(generated[0], skip_special_tokens=True)
+    
+    def _eval_sequence_identity(self, model, tokenizer, task: BioTask) -> float:
+        """Evaluate exact sequence identity."""
+        prompt = task.prompt
+        if task.context:
+            prompt += f" {task.context}"
+        
+        output = self._get_model_output(model, tokenizer, prompt)
+        
+        if task.expected_answer is None:
+            return 0.5  # Default if no expected answer
+        
+        # Extract sequence from output
+        output_seq = self._extract_sequence(output, task.task_type)
+        expected = task.expected_answer.strip().upper()
+        
+        if not output_seq or not expected:
+            return 0.0
+        
+        # Compute identity
+        matches = sum(1 for a, b in zip(output_seq, expected) if a == b)
+        length = max(len(output_seq), len(expected))
+        
+        return matches / length if length > 0 else 0.0
+    
+    def _eval_sequence_similarity(self, model, tokenizer, task: BioTask) -> float:
+        """Evaluate sequence similarity using multiple metrics."""
+        prompt = task.prompt
+        if task.context:
+            prompt += f" {task.context}"
+        
+        output = self._get_model_output(model, tokenizer, prompt)
+        
+        if task.expected_answer is None:
+            return 0.5
+        
+        output_seq = self._extract_sequence(output, task.task_type)
+        expected = task.expected_answer.strip().upper()
+        
+        if not output_seq or not expected:
+            return 0.0
+        
+        # SequenceMatcher ratio
+        sm = SequenceMatcher(None, output_seq, expected)
+        similarity = sm.ratio()
+        
+        # Also compute local alignment score (simplified)
+        # Could use Bio.pairwise2 or biopython for full alignment
+        
+        return similarity
+    
+    def _eval_contains_substring(self, model, tokenizer, task: BioTask) -> float:
+        """Check if output contains expected motif."""
+        prompt = task.prompt
+        if task.context:
+            prompt += f" {task.context}"
+        
+        output = self._get_model_output(model, tokenizer, prompt)
+        
+        if task.expected_answer is None:
+            return 0.5
+        
+        expected = task.expected_answer.strip().upper()
+        output_seq = self._extract_sequence(output, task.task_type)
+        
+        if expected in output_seq:
+            return 1.0
+        
+        # Partial match
+        for i in range(len(expected) - 2):
+            sub = expected[i:i+3]
+            if sub in output_seq:
+                return 0.3
+        
+        return 0.0
+    
+    def _eval_exact_match(self, model, tokenizer, task: BioTask) -> float:
+        """Exact match evaluation."""
+        prompt = task.prompt
+        if task.context:
+            prompt += f" {task.context}"
+        
+        output = self._get_model_output(model, tokenizer, prompt)
+        
+        if task.expected_answer is None:
+            return 0.5
+        
+        # Extract answer from output
+        output_answer = self._extract_answer(output)
+        expected = task.expected_answer.strip()
+        
+        if output_answer == expected:
+            return 1.0
+        
+        # Numeric approximate match
+        try:
+            output_num = float(output_answer)
+            expected_num = float(expected)
+            if abs(output_num - expected_num) < 1:
+                return 0.5
+        except (ValueError, TypeError):
+            pass
+        
+        return 0.0
+    
+    def _eval_perplexity(self, model, tokenizer, task: BioTask) -> float:
+        """Evaluate perplexity on a sequence."""
+        if task.target is None:
+            return 0.5
+        
+        text = task.target
+        input_ids = tokenizer.encode(text, return_tensors="pt", max_length=self.max_length, truncation=True)
+        input_ids = input_ids.to(self.device)
+        
+        with torch.no_grad():
+            outputs = model(input_ids, labels=input_ids)
+            loss = outputs.loss
+        
+        perplexity = torch.exp(loss).item()
+        
+        # Convert to score (lower perplexity = higher score)
+        # Typical perplexity for protein LMs is 5-20
+        score = 1.0 / (1.0 + perplexity / 10.0)
+        
+        return score
+    
+    def _eval_rna_structure(self, model, tokenizer, task: BioTask) -> float:
+        """
+        Evaluate RNA structure prediction.
+        Uses simplified dot-bracket notation comparison.
+        """
+        prompt = task.prompt
+        if task.context:
+            prompt += f" {task.context}"
+        
+        output = self._get_model_output(model, tokenizer, prompt)
+        
+        # Extract predicted structure (dot-bracket notation)
+        predicted = self._extract_structure(output)
+        
+        if not predicted:
+            return 0.0
+        
+        # For generated tasks without expected structure, just check validity
+        if task.expected_answer is None:
+            # Check if dot-bracket is balanced
+            balance = 0
+            valid = True
+            for c in predicted:
+                if c == '(':
+                    balance += 1
+                elif c == ')':
+                    balance -= 1
+                    if balance < 0:
+                        valid = False
+            
+            if valid and balance == 0:
+                return 0.5
+            return 0.0
+        
+        expected = task.expected_answer
+        
+        # Compare structures
+        matches = sum(1 for a, b in zip(predicted, expected) if a == b)
+        return matches / max(len(predicted), len(expected))
+    
+    def _extract_sequence(self, text: str, seq_type: str) -> str:
+        """Extract biological sequence from model output."""
+        # Remove special tokens and whitespace
+        text = text.replace("<mask>", "").replace("[MASK]", "")
+        text = text.replace("<s>", "").replace("</s>", "")
+        text = text.replace("[CLS]", "").replace("[SEP]", "")
+        
+        # For proteins, look for uppercase amino acid sequences
+        if seq_type == "protein":
+            pattern = re.compile(r'[ACDEFGHIKLMNPQRSTVWY]+')
+            matches = pattern.findall(text.upper())
+            if matches:
+                return max(matches, key=len)
+            return text.upper()
+        
+        # For DNA
+        elif seq_type == "dna":
+            pattern = re.compile(r'[ACGT]+')
+            matches = pattern.findall(text.upper())
+            if matches:
+                return max(matches, key=len)
+            return text.upper().replace('U', 'T')
+        
+        # For RNA
+        elif seq_type == "rna":
+            pattern = re.compile(r'[ACGU]+')
+            matches = pattern.findall(text.upper())
+            if matches:
+                return max(matches, key=len)
+            return text.upper().replace('T', 'U')
+        
+        return text.upper().strip()
+    
+    def _extract_answer(self, text: str) -> str:
+        """Extract a short answer from model output."""
+        # Try to find a number
+        numbers = re.findall(r'-?\d+', text)
+        if numbers:
+            return numbers[-1]  # Last number is often the answer
+        
+        # Or take the last non-empty line
+        lines = [l.strip() for l in text.split('\n') if l.strip()]
+        if lines:
+            return lines[-1]
+        
+        return text.strip()
+    
+    def _extract_structure(self, text: str) -> str:
+        """Extract dot-bracket RNA structure notation."""
+        pattern = re.compile(r'[\(\)\.]+')
+        matches = pattern.findall(text)
+        if matches:
+            return max(matches, key=len)
+        return ""