Add final_sequence_decoder.py

src/final_sequence_decoder.py

@@ -0,0 +1,263 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+import esm
+from tqdm import tqdm
+import os
+from datetime import datetime
+
+CANONICAL_AAS = list("ACDEFGHIKLMNPQRSTVWY")
+
+class EmbeddingToSequenceConverter:
+    """
+    Decode contextual ESM2 hidden states to amino-acid sequences via the model's LM head.
+    Accepts [L, 1280] or [B, L, 1280] tensors (L≈50 in your pipeline).
+    """
+
+    def __init__(self, device="cuda"):
+        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
+        print("Loading ESM model for sequence decoding...")
+        self.model, self.alphabet = esm.pretrained.esm2_t33_650M_UR50D()
+        self.model.eval().to(self.device)
+        self.aa_list = CANONICAL_AAS
+        self.aa_token_ids = torch.tensor(
+            [self.alphabet.get_idx(a) for a in self.aa_list],
+            device=self.device, dtype=torch.long
+        )
+        print("✓ ESM model loaded for sequence decoding")
+    
+    @torch.inference_mode()
+    def _logits_from_hidden(self, hidden):
+        # hidden: [L, D] or [B, L, D]; project exactly as ESM-2 does (LayerNorm → LM head)
+        if hidden.dim() == 2:
+            hidden = hidden.unsqueeze(0)
+        hidden = hidden.to(self.device)
+        # match model dtype to avoid dtype mismatches under autocast
+        hidden = hidden.to(self.model.lm_head.weight.dtype)
+        if hasattr(self.model, "emb_layer_norm_after"):
+            hidden = self.model.emb_layer_norm_after(hidden)
+        logits_full = self.model.lm_head(hidden)                        # [B, L, |V|]
+        logits_20  = logits_full.index_select(-1, self.aa_token_ids)    # [B, L, 20]
+        return logits_20
+
+    @torch.inference_mode()
+    def embedding_to_sequence(self, embedding, method="diverse", temperature=0.8, top_p=0.9, top_k=0, seed=None, return_conf=False):
+        logits = self._logits_from_hidden(embedding)                    # [1, L, 20]
+        if method in ("nearest", "nearest_neighbor"):
+            idx = logits.argmax(-1)[0]
+            probs = logits.softmax(-1)[0]
+        else:
+            if seed is not None:
+                torch.manual_seed(seed)
+            if temperature and temperature > 0:
+                logits = logits / temperature
+            probs = logits.softmax(-1)[0]                               # [L, 20]
+            V = probs.size(-1)
+            if top_k and top_k < V:
+                kth = torch.topk(probs, top_k, dim=-1).values[..., -1:]
+                probs = torch.where(probs >= kth, probs, torch.zeros_like(probs))
+                probs = probs / probs.sum(-1, keepdim=True).clamp_min(1e-12)
+            if top_p and 0 < top_p < 1:
+                sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
+                cum = sorted_probs.cumsum(-1)
+                mask = cum > top_p
+                mask[..., 0] = False
+                sorted_probs = sorted_probs.masked_fill(mask, 0)
+                sorted_probs = sorted_probs / sorted_probs.sum(-1, keepdim=True).clamp_min(1e-12)
+                samples = torch.multinomial(sorted_probs, 1).squeeze(-1)
+                idx = sorted_idx.gather(-1, samples.unsqueeze(-1)).squeeze(-1)
+            else:
+                idx = torch.multinomial(probs, 1).squeeze(-1)
+        seq = "".join(self.aa_list[i] for i in idx.tolist())
+        if return_conf:
+            conf = probs.max(-1).values.mean().item()   # avg per-pos max prob
+            return seq, conf
+        return seq
+
+    @torch.inference_mode()
+    def batch_embedding_to_sequences(self, embeddings, method="diverse", temperature=0.8, top_p=0.9, top_k=0, seed=None, return_conf=False, max_tokens=100_000):
+        if embeddings.dim() == 2:
+            return [self.embedding_to_sequence(embeddings, method, temperature, top_p, top_k, seed, return_conf)]
+        B, L, V = embeddings.shape
+        if seed is not None:
+            torch.manual_seed(seed)
+        # Batched logits to avoid OOM
+        logits = []
+        start = 0
+        while start < B:
+            chunk_bs = max(1, min(B - start, max_tokens // L))
+            logits.append(self._logits_from_hidden(embeddings[start:start+chunk_bs]))
+            start += chunk_bs
+        logits = torch.cat(logits, dim=0)                              # [B, L, 20]
+        if method in ("nearest", "nearest_neighbor"):
+            idx = logits.argmax(-1)                                     # [B, L]
+            probs = logits.softmax(-1)
+        else:
+            if temperature and temperature > 0:
+                logits = logits / temperature
+            probs = logits.softmax(-1)                                  # [B, L, 20]
+            B, L, V = probs.shape
+            if top_k and top_k < V:
+                kth = torch.topk(probs, top_k, dim=-1).values[..., -1:].expand_as(probs)
+                probs = torch.where(probs >= kth, probs, torch.zeros_like(probs))
+                probs = probs / probs.sum(-1, keepdim=True).clamp_min(1e-12)
+            if top_p and 0 < top_p < 1:
+                flat = probs.view(-1, V)
+                sorted_probs, sorted_idx = torch.sort(flat, descending=True, dim=-1)
+                cum = sorted_probs.cumsum(-1)
+                mask = cum > top_p
+                mask[:, 0] = False
+                sorted_probs = sorted_probs.masked_fill(mask, 0)
+                sorted_probs = sorted_probs / sorted_probs.sum(-1, keepdim=True).clamp_min(1e-12)
+                samples = torch.multinomial(sorted_probs, 1)            # [B*L, 1]
+                idx = sorted_idx.gather(-1, samples).view(B, L)         # [B, L]
+            else:
+                idx = torch.multinomial(probs.view(-1, V), 1).view(B, L)
+        seqs = ["".join(self.aa_list[i] for i in row.tolist()) for row in idx]
+        if return_conf:
+            conf = probs.max(-1).values.mean(-1).tolist()               # [B]
+            return list(zip(seqs, conf))
+        return seqs
+    def validate_sequence(self, s):
+        return all(a in set(self.aa_list) for a in s)
+    
+    def filter_valid_sequences(self, sequences):
+        valid = []
+        for seq in sequences:
+            if self.validate_sequence(seq):
+                valid.append(seq)
+            else:
+                print(f"Warning: Invalid sequence found: {seq}")
+        return valid
+
+def main():
+    """
+    Decode all CFG-generated peptide embeddings to sequences and analyze distribution.
+    Uses the best trained model (loss: 0.017183, step: 53).
+    """
+    print("=== CFG-Generated Peptide Sequence Decoder (Best Model) ===")
+    
+    # Initialize converter
+    converter = EmbeddingToSequenceConverter()
+    
+    # Get today's date for filename
+    today = datetime.now().strftime('%Y%m%d')
+    
+    # Load all CFG-generated embeddings (using best model)
+    cfg_files = {
+        'No CFG (0.0)': f'/data2/edwardsun/generated_samples/generated_amps_best_model_no_cfg_{today}.pt',
+        'Weak CFG (3.0)': f'/data2/edwardsun/generated_samples/generated_amps_best_model_weak_cfg_{today}.pt', 
+        'Strong CFG (7.5)': f'/data2/edwardsun/generated_samples/generated_amps_best_model_strong_cfg_{today}.pt',
+        'Very Strong CFG (15.0)': f'/data2/edwardsun/generated_samples/generated_amps_best_model_very_strong_cfg_{today}.pt'
+    }
+    
+    all_results = {}
+    
+    for cfg_name, file_path in cfg_files.items():
+        print(f"\n{'='*50}")
+        print(f"Processing {cfg_name}...")
+        print(f"Loading: {file_path}")
+        
+        try:
+            # Load embeddings
+            embeddings = torch.load(file_path, map_location='cpu')
+            print(f"✓ Loaded {len(embeddings)} embeddings, shape: {embeddings.shape}")
+            
+            # Decode to sequences using diverse method
+            print(f"Decoding sequences...")
+            sequences = converter.batch_embedding_to_sequences(embeddings, method='diverse', temperature=0.5)
+            
+            # Filter valid sequences
+            valid_sequences = converter.filter_valid_sequences(sequences)
+            print(f"✓ Valid sequences: {len(valid_sequences)}/{len(sequences)}")
+            
+            # Store results
+            all_results[cfg_name] = {
+                'sequences': valid_sequences,
+                'total': len(sequences),
+                'valid': len(valid_sequences)
+            }
+            
+            # Show sample sequences
+            print(f"\nSample sequences ({cfg_name}):")
+            for i, seq in enumerate(valid_sequences[:5]):
+                print(f"  {i+1}: {seq}")
+            
+        except Exception as e:
+            print(f"❌ Error processing {file_path}: {e}")
+            all_results[cfg_name] = {'sequences': [], 'total': 0, 'valid': 0}
+    
+    # Analysis and comparison
+    print(f"\n{'='*60}")
+    print("CFG ANALYSIS SUMMARY")
+    print(f"{'='*60}")
+    
+    for cfg_name, results in all_results.items():
+        sequences = results['sequences']
+        if sequences:
+            # Calculate sequence statistics
+            lengths = [len(seq) for seq in sequences]
+            avg_length = np.mean(lengths)
+            std_length = np.std(lengths)
+            
+            # Calculate amino acid composition
+            all_aas = ''.join(sequences)
+            aa_counts = {}
+            for aa in 'ACDEFGHIKLMNPQRSTVWY':
+                aa_counts[aa] = all_aas.count(aa)
+            
+            # Calculate diversity (unique sequences)
+            unique_sequences = len(set(sequences))
+            diversity_ratio = unique_sequences / len(sequences)
+            
+            print(f"\n{cfg_name}:")
+            print(f"  Total sequences: {results['total']}")
+            print(f"  Valid sequences: {results['valid']}")
+            print(f"  Unique sequences: {unique_sequences}")
+            print(f"  Diversity ratio: {diversity_ratio:.3f}")
+            print(f"  Avg length: {avg_length:.1f} ± {std_length:.1f}")
+            print(f"  Length range: {min(lengths)}-{max(lengths)}")
+            
+            # Show top amino acids
+            sorted_aas = sorted(aa_counts.items(), key=lambda x: x[1], reverse=True)
+            print(f"  Top 5 AAs: {', '.join([f'{aa}({count})' for aa, count in sorted_aas[:5]])}")
+            
+            # Create output directory if it doesn't exist
+            output_dir = '/data2/edwardsun/decoded_sequences'
+            os.makedirs(output_dir, exist_ok=True)
+            
+            # Save sequences to file with date
+            output_file = os.path.join(output_dir, f"decoded_sequences_{cfg_name.lower().replace(' ', '_').replace('(', '').replace(')', '').replace('.', '')}_{today}.txt")
+            with open(output_file, 'w') as f:
+                f.write(f"# Decoded sequences from {cfg_name}\n")
+                f.write(f"# Total: {results['total']}, Valid: {results['valid']}, Unique: {unique_sequences}\n")
+                f.write(f"# Generated from best model (loss: 0.017183, step: 53)\n\n")
+                for i, seq in enumerate(sequences):
+                    f.write(f"seq_{i+1:03d}\t{seq}\n")
+            print(f"  ✓ Saved to: {output_file}")
+    
+    # Overall comparison
+    print(f"\n{'='*60}")
+    print("OVERALL COMPARISON")
+    print(f"{'='*60}")
+    
+    cfg_names = list(all_results.keys())
+    valid_counts = [all_results[name]['valid'] for name in cfg_names]
+    unique_counts = [len(set(all_results[name]['sequences'])) for name in cfg_names]
+    
+    print(f"Valid sequences: {dict(zip(cfg_names, valid_counts))}")
+    print(f"Unique sequences: {dict(zip(cfg_names, unique_counts))}")
+    
+    # Find most diverse and most similar
+    if all(valid_counts):
+        diversity_ratios = [unique_counts[i]/valid_counts[i] for i in range(len(valid_counts))]
+        most_diverse = cfg_names[diversity_ratios.index(max(diversity_ratios))]
+        least_diverse = cfg_names[diversity_ratios.index(min(diversity_ratios))]
+        
+        print(f"\nMost diverse: {most_diverse} (ratio: {max(diversity_ratios):.3f})")
+        print(f"Least diverse: {least_diverse} (ratio: {min(diversity_ratios):.3f})")
+    
+    print(f"\n✓ Decoding complete! Check the output files for detailed sequences.")
+
+if __name__ == "__main__":
+    main()