models.py

from abc import ABC, abstractmethod
import torch
from progen2 import progen2generate
from transformers import PreTrainedTokenizerFast
import torch.nn.functional as F
import os
import logging
from typing import Tuple, Optional
import pandas as pd
import time

def read_fasta(file_path):
    """
    Reads a FASTA file and returns the sequence.
    """
    with open(file_path, 'r') as file:
        lines = file.readlines()
    
    sequence_lines = [line.strip() for line in lines if not line.startswith('>')]
    sequence = ''.join(sequence_lines)
    return sequence

class ProteinLanguageModel(ABC):
    # Standard amino acid ordering that all models should use
    standard_aa_order = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 
                        'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']

    @abstractmethod
    def generate(self, context: str, max_new_tokens: int, temperature: float = 1.0):
        """Generate a sequence from the given context."""
        pass

    @abstractmethod
    def get_logits(self, input_ids: torch.Tensor, past_key_values=None, use_cache=False):
        """Get logits for the input sequence.
        
        Args:
            input_ids: Input token IDs
            past_key_values: Optional past key values for faster inference
            use_cache: Whether to return updated key values
            
        Returns:
            Logits tensor or tuple of (logits, past_key_values) if use_cache=True
        """
        pass

    @abstractmethod
    def get_tokenizer(self):
        """Get the model's tokenizer."""
        pass

class ProGen2Model(ProteinLanguageModel):
    def __init__(self, size='small', device='cuda:0'):
        self.model, self.tokenizer, self.device = progen2generate.load_model(size=size)
        self.model = self.model.to(device)
        self.device = device

        # Create mapping from ProGen2 token IDs to standard ordering
        self.model_to_standard = {}
        for idx, aa in enumerate(self.standard_aa_order):
            token_id = self.tokenizer.token_to_id(aa)
            if token_id is not None:
                self.model_to_standard[token_id] = idx
        # Add special tokens
        self.model_to_standard[self.tokenizer.token_to_id('<|endoftext|>')] = 20  # End token
        self.model_to_standard[self.tokenizer.token_to_id('<|pad|>')] = 21  # Pad token

    def encode(self, text):
        """Convert text to token IDs."""
        start_time = time.time()
        encoding = self.tokenizer.encode(text)
        tokens = encoding.ids  # Extract the token IDs from the Encoding object
        return tokens

    def decode(self, token_ids):
        """Convert token IDs back to text."""
        start_time = time.time()
        # First decode to text
        text = self.tokenizer.decode(token_ids)
        # Remove all special tokens
        special_tokens = ["[SEP]", "[PAD]", "[CLS]", "[UNK]", "[MASK]"]
        for token in special_tokens:
            text = text.replace(token, "")
        # Remove all whitespace, including spaces between amino acids
        text = "".join(text.split())
        return text

    def generate(self, context: str, max_new_tokens: int, temperature: float = 1.0, num_return_sequences: int = 1):
        """Generate sequences and return completions, likelihoods, and timing info."""
        start_time = time.time()

        if torch.cuda.is_available():
            torch.cuda.synchronize()

        memory_before = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0
        
        # Call generate_sequence which now returns timing info
        completions, likelihoods = progen2generate.generate_sequence(
            self.model,
            self.tokenizer,
            self.device,
            context,
            temp=temperature,
            max_new_tokens=max_new_tokens,
            num_return_sequences=num_return_sequences
        )

        if torch.cuda.is_available():
            torch.cuda.synchronize()

        memory_after = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0
        memory_diff = memory_after - memory_before

        end_time = time.time()
        # Return including timing info
        return completions, likelihoods

    def get_logits(self, input_ids: torch.Tensor, past_key_values=None, use_cache=False):
        start_time = time.time()
        
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            
        # Call the model with appropriate caching parameters
        if use_cache:
            outputs = self.model(input_ids, past_key_values=past_key_values, use_cache=True)
            # If this is a huggingface model, it will return an object with logits and past_key_values
            if hasattr(outputs, 'past_key_values'):
                logits = outputs.logits
                past_key_values = outputs.past_key_values
            else:
                # If it's a tuple of (logits, past_key_values)
                if isinstance(outputs, tuple) and len(outputs) == 2:
                    logits, past_key_values = outputs
                else:
                    logits = outputs
        else:
            outputs = self.model(input_ids)
            logits = outputs.logits if hasattr(outputs, 'logits') else outputs
        
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            
        end_time = time.time()
        
        if use_cache:
            return logits, past_key_values
        return logits

    def get_tokenizer(self):
        return self.tokenizer

    def id_to_token(self, token_id):
        """Convert token ID to amino acid or special token."""
        if token_id in self.model_to_standard:
            # If it's an amino acid, find it in the standard ordering
            standard_idx = self.model_to_standard[token_id]
            if standard_idx < 20:  # If it's one of the 20 amino acids
                return self.standard_aa_order[standard_idx]
        # For special tokens or unknown IDs, use the tokenizer's conversion
        return self.tokenizer.id_to_token(token_id)

class TranceptionModel(ProteinLanguageModel):
    def __init__(self, size='Small', device='cuda:0'):
        self.device = device
        # Set up custom cache directory
        cache_dir = "./model_cache/tranception"
        os.makedirs(cache_dir, exist_ok=True)
        
        self.tokenizer = PreTrainedTokenizerFast(
            tokenizer_file="./Tranception/tranception/utils/tokenizers/Basic_tokenizer",
            unk_token="[UNK]",
            sep_token="[SEP]",
            pad_token="[PAD]",
            cls_token="[CLS]",
            mask_token="[MASK]"
        )
        
        # Create mapping between token IDs and amino acids
        self.aa_to_id = {
            aa: self.tokenizer.convert_tokens_to_ids(aa) 
            for aa in self.standard_aa_order
        }
        self.id_to_aa = {v: k for k, v in self.aa_to_id.items()}

        # Create mapping from Tranception token IDs to standard ordering
        self.model_to_standard = {self.aa_to_id[aa]: idx for idx, aa in enumerate(self.standard_aa_order)}
        # Add special tokens
        self.model_to_standard[self.tokenizer.sep_token_id] = 20  # End token
        self.model_to_standard[self.tokenizer.pad_token_id] = 21  # Pad token

        # Store stop tokens for easy access
        self.stop_tokens = [self.tokenizer.sep_token_id, self.tokenizer.pad_token_id]

        # Load the appropriate model size
        if size == 'Small':
            model_name = "PascalNotin/Tranception_Small"
        elif size == 'Medium':
            model_name = "PascalNotin/Tranception_Medium"
        else:
            model_name = "PascalNotin/Tranception_Large"
        
        # Initialize the model with the custom configuration
        self.model = TranceptionLMHeadModel.from_pretrained(
            model_name,
            cache_dir=cache_dir,
            local_files_only=False
        )
        
        self.model = self.model.to(device)
        self.model.config.tokenizer = self.tokenizer

    def encode(self, text):
        """Convert text to token IDs."""
        return self.tokenizer.encode(text)

    def decode(self, token_ids):
        """Convert token IDs back to text."""
        # First decode to text
        text = self.tokenizer.decode(token_ids)
        # Remove all special tokens
        special_tokens = ["[SEP]", "[PAD]", "[CLS]", "[UNK]", "[MASK]"]
        for token in special_tokens:
            text = text.replace(token, "")
        # Remove all whitespace, including spaces between amino acids
        text = "".join(text.split())
        return text

    def generate(self, context: str, max_new_tokens: int, temperature: float = 1.0):
        input_ids = torch.tensor([self.tokenizer.encode(context)], device=self.device)
        
        generated_sequence = []
        all_log_probs = []
        
        for _ in range(max_new_tokens):
            with torch.no_grad():
                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=torch.ones_like(input_ids),
                    return_dict=True
                )
                logits = outputs.logits[:, -1, :] / temperature
                
                # Create a mask for amino acid tokens only
                aa_mask = torch.zeros_like(logits, dtype=torch.bool)
                for aa_id in self.id_to_aa.keys():
                    aa_mask[0, aa_id] = True
                
                # Mask out everything except amino acids
                logits = logits.masked_fill(~aa_mask, float('-inf'))
                
                # Compute log probabilities
                log_probs = F.log_softmax(logits, dim=-1)
                probs = torch.exp(log_probs)
                
                # Sample from the distribution
                next_token = torch.multinomial(probs, num_samples=1)
                
                # Store log probabilities as tensor
                all_log_probs.append(log_probs)
                
                # Append to generated sequence
                generated_sequence.append(next_token.item())
                input_ids = torch.cat([input_ids, next_token], dim=-1)
        
        complete_sequence = self.decode(generated_sequence)
        # Stack all log probabilities into a single tensor
        all_log_probs = torch.stack(all_log_probs, dim=0).squeeze(1)
        return complete_sequence, all_log_probs

    def get_logits(self, input_ids: torch.Tensor, past_key_values=None, use_cache=False):
        start_time = time.time()
        
        if use_cache:
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=torch.ones_like(input_ids),
                past_key_values=past_key_values,
                use_cache=True,
                return_dict=True
            )
            # Tranception model should return past_key_values in the output
            logits = outputs.logits
            past_key_values = outputs.past_key_values
            
            end_time = time.time()
            
            return logits, past_key_values
        else:
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=torch.ones_like(input_ids),
                return_dict=True
            )
            
            end_time = time.time()
            
            return outputs.logits

    def get_tokenizer(self):
        return self.tokenizer

    def id_to_token(self, token_id):
        """Convert token ID to amino acid or special token."""
        if token_id in self.id_to_aa:
            return self.id_to_aa[token_id]
        # For special tokens, use the tokenizer's conversion
        return self.tokenizer.convert_ids_to_tokens([token_id])[0] 

    def handle_stop_token(self, draft_token: int, target_probs: torch.Tensor, p_x: float, eta: float) -> Tuple[bool, Optional[int]]:
        """Handle special stop token logic for Tranception."""
        if draft_token in self.stop_tokens:
            # Never accept stop tokens, always sample from target distribution excluding stop tokens
            target_probs = target_probs.clone()  # Create a copy to modify
            target_probs[self.stop_tokens] = 0
            target_probs = target_probs / target_probs.sum()
            correction_token = torch.multinomial(target_probs, 1).item()
            return False, correction_token
        return None, None  # Not a stop token, use normal coupling 

class MLPDraftModel(ProteinLanguageModel):
    """A model that uses pre-computed MLP score matrices for generation."""
    def __init__(self, protein_name: str, device: str = 'cuda'):
        self.device = device
        self.protein_name = protein_name
        
        # Load score matrix
        if protein_name == "GFP":
            score_matrix_path = "./evaluation/MLP_GFP_logps.csv"
        elif protein_name == "RBP1":
            score_matrix_path = "./evaluation/MLP_RBP1_logps.csv"
        elif protein_name == "F7YBW8":
            score_matrix_path = "./evaluation/MLP_F7YBW8_logps.csv"
        else:
            raise ValueError(f"Unknown protein: {protein_name}")
            
        self.score_matrix = pd.read_csv(score_matrix_path, index_col=0)
        
        # Create amino acid to index mapping using standard order
        self.aa_to_idx = {aa: i for i, aa in enumerate(self.standard_aa_order)}
        self.idx_to_aa = {i: aa for aa, i in self.aa_to_idx.items()}
        
        # Add special tokens
        self.aa_to_idx.update({'1': 20, '2': 21})  # Start and stop tokens
        self.idx_to_aa.update({20: '1', 21: '2'})
        
        # Load wildtype sequence for context
        if protein_name == "GFP":
            self.wildtype = read_fasta(f'./data/avGFP.fasta')
        else:
            self.wildtype = read_fasta(f'./data/{protein_name}.fasta')
            
    def encode(self, sequence: str) -> list:
        """Convert an amino acid sequence to token IDs."""
        return [self.aa_to_idx.get(aa, -1) for aa in sequence]
        
    def decode(self, tokens: list) -> str:
        """Convert token IDs back to amino acid sequence."""
        return ''.join(self.idx_to_aa.get(t, 'X') for t in tokens)
        
    def get_tokenizer(self):
        """Return a simple tokenizer that works with single amino acids."""
        class SimpleTokenizer:
            def __init__(self_, model):
                self_.model = model
                
            def encode(self_, sequence):
                return self_.model.encode(sequence)
            
            def decode(self_, tokens):
                if isinstance(tokens, torch.Tensor):
                    tokens = tokens.tolist()
                return self_.model.decode(tokens)
                
        return SimpleTokenizer(self)
    
    def get_logits(self, input_ids: torch.Tensor):
        """Get logits for the next token."""
        if isinstance(input_ids, torch.Tensor):
            input_ids = input_ids.tolist()
            
        current_pos = len(input_ids[0]) - 1
        if current_pos >= len(self.score_matrix.columns):
            # If beyond score matrix, return uniform distribution
            logits = torch.zeros((1, len(self.aa_to_idx)), device=self.device)
            return logits
            
        # Get scores for this position
        pos_scores = self.score_matrix[self.score_matrix.columns[current_pos]].values
        
        # Create full logits including special tokens
        full_logits = torch.full((1, len(self.aa_to_idx)), float('-inf'), device=self.device)
        full_logits[0, :20] = torch.tensor(pos_scores, device=self.device)  # Standard AAs
        full_logits[0, 21] = 0.0  # Allow stop token with neutral probability
        
        return full_logits
    
    def generate(self, context: str, max_new_tokens: int, temperature: float = 1.0):
        """Generate a sequence using the MLP score matrix."""
        sequence = list(context)
        tokens = self.encode(context)
        
        for pos in range(len(context), len(context) + max_new_tokens):
            if pos >= len(self.score_matrix.columns):
                break
                
            # Get logits and apply temperature
            logits = self.get_logits(torch.tensor([tokens], device=self.device))
            if temperature != 1.0:
                logits = logits / temperature
            
            # Convert to probabilities
            probs = torch.softmax(logits, dim=-1)
            
            # Sample next token
            next_token = torch.multinomial(probs[0], 1).item()
            next_aa = self.idx_to_aa[next_token]
            
            sequence.append(next_aa)
            tokens.append(next_token)
            
            if next_token == 21:  # Stop token
                break
        
        return ''.join(sequence), tokens 

class FineTunedProGen2Model(ProteinLanguageModel):
    def __init__(self, model_path, device='cuda:0'):
        from progen2.sample import create_model, create_tokenizer_custom
        
        # Load tokenizer
        tokenizer_path = os.path.join(model_path, "tokenizer.json")
        if os.path.exists(tokenizer_path):
            self.tokenizer = create_tokenizer_custom(file=tokenizer_path)
        else:
            raise ValueError(f"Tokenizer not found at {tokenizer_path}")
        
        # Load model
        self.model = create_model(ckpt=model_path, fp16=False)
        self.model = self.model.to(device)
        self.device = device

        # Create mapping from ProGen2 token IDs to standard ordering
        self.model_to_standard = {}
        for idx, aa in enumerate(self.standard_aa_order):
            token_id = self.tokenizer.token_to_id(aa)
            if token_id is not None:
                self.model_to_standard[token_id] = idx
        # Add special tokens
        self.model_to_standard[self.tokenizer.token_to_id('<|endoftext|>')] = 20  # End token
        self.model_to_standard[self.tokenizer.token_to_id('<|pad|>')] = 21  # Pad token

    def encode(self, text):
        """Convert text to token IDs."""
        encoding = self.tokenizer.encode(text)
        return encoding.ids  # Extract the token IDs from the Encoding object

    def decode(self, token_ids):
        """Convert token IDs back to text."""
        # First decode to text
        text = self.tokenizer.decode(token_ids)
        # Remove all special tokens
        special_tokens = ["[SEP]", "[PAD]", "[CLS]", "[UNK]", "[MASK]"]
        for token in special_tokens:
            text = text.replace(token, "")
        # Remove all whitespace, including spaces between amino acids
        text = "".join(text.split())
        return text

    def generate(self, context: str, max_new_tokens: int, temperature: float = 1.0):
        from progen2 import progen2generate
        
        completions, likelihoods = progen2generate.generate_sequence(
            self.model,
            self.tokenizer,
            self.device,
            context,
            temp=temperature,
            max_new_tokens=max_new_tokens
        )
        return completions[0], likelihoods

    def get_logits(self, input_ids: torch.Tensor, past_key_values=None, use_cache=False):
        outputs = self.model(input_ids)
        return outputs.logits

    def get_tokenizer(self):
        return self.tokenizer

    def id_to_token(self, token_id):
        """Convert token ID to amino acid or special token."""
        if token_id in self.model_to_standard:
            # If it's an amino acid, find it in the standard ordering
            standard_idx = self.model_to_standard[token_id]
            if standard_idx < 20:  # If it's one of the 20 amino acids
                return self.standard_aa_order[standard_idx]
        # For special tokens or unknown IDs, use the tokenizer's conversion
        return self.tokenizer.id_to_token(token_id)