#!/usr/bin/env python3 """ V5.1-FINAL: Corrected Negative Generation FIXES from v3: 1. Sialic acid detection now uses WURCS pattern 'a2122h' (Neu5Ac) and 'a2212h' (Neu5Gc) 2. All token pattern matching updated for WURCS format 3. Rule5 (Sia extension) should now work correctly OUTPUT: 100K rule-based negatives with proper biological patterns """ import sys import pickle import json import random import re import numpy as np from copy import deepcopy from typing import Dict, List, Optional, Tuple from tqdm import tqdm import logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Paths BASE = "/work/ratul1/supantha/glycan-SD-VS/bert_training_v3/v3.1_cluster_training" BPE_VOCAB_PATH = f"{BASE}/data/bpe_vocabulary.json" POSITIVES_PATH = f"{BASE}/bert_v5.1_contrastive/data/fully_resolved_161k.pkl" OUTPUT_PATH = f"{BASE}/bert_v5.1_contrastive/data/hard_negatives_100k_v4_FINAL.pkl" N_SAMPLES = 100000 class WURCSPatterns: """WURCS-specific patterns for monosaccharides and linkages.""" # Monosaccharide patterns in WURCS SIALIC_ACID = ['a2122h', 'a2212h'] # Neu5Ac, Neu5Gc IDURONIC_ACID = ['a1221m'] # IdoA GLUCURONIC_ACID = ['a2112h'] # GlcA MANNOSE = ['a2112m'] # Man GLUCOSE = ['a2122m'] # Glc GALACTOSE = ['a2112a'] # Gal # Linkage patterns ALPHA_LINKS = ['-1a', 'a-'] # Alpha anomeric BETA_LINKS = ['-1b', 'b-'] # Beta anomeric class UniversalBiologicalRulesV4: """ V4: Proper WURCS pattern detection for all 8 rules. """ def __init__(self, vocab_path: str, sequences: List[Dict]): logger.info(f"Loading BPE vocabulary from {vocab_path}") with open(vocab_path, 'r') as f: vocab = json.load(f) self.token_to_id = vocab['token_to_id'] self.id_to_token = {v: k for k, v in self.token_to_id.items()} self.vocab_size = len(self.token_to_id) self.sequences = sequences # Categorize tokens by WURCS patterns self._categorize_tokens() def _categorize_tokens(self): """Categorize vocabulary tokens by WURCS patterns.""" self.sialic_tokens = [] self.idoa_tokens = [] self.alpha_linkage_tokens = [] self.beta_linkage_tokens = [] self.monosaccharide_tokens = [] for token, tid in self.token_to_id.items(): # Sialic acid (Neu5Ac = a2122h, Neu5Gc = a2212h) if any(pattern in token for pattern in WURCSPatterns.SIALIC_ACID): self.sialic_tokens.append(tid) # Iduronic acid (IdoA = a1221m) if any(pattern in token for pattern in WURCSPatterns.IDURONIC_ACID): self.idoa_tokens.append(tid) # Alpha linkages (-1a or ends with 'a') if any(pattern in token for pattern in WURCSPatterns.ALPHA_LINKS): self.alpha_linkage_tokens.append(tid) # Beta linkages (-1b or ends with 'b') if any(pattern in token for pattern in WURCSPatterns.BETA_LINKS): self.beta_linkage_tokens.append(tid) # General monosaccharides (start with 'a' followed by digits) if re.match(r'^a\d+[a-z]', token): self.monosaccharide_tokens.append(tid) logger.info(f"WURCS Token Categories:") logger.info(f" Sialic acid (a2122h/a2212h): {len(self.sialic_tokens)} tokens") logger.info(f" IdoA (a1221m): {len(self.idoa_tokens)} tokens") logger.info(f" Alpha linkages: {len(self.alpha_linkage_tokens)} tokens") logger.info(f" Beta linkages: {len(self.beta_linkage_tokens)} tokens") logger.info(f" Monosaccharides: {len(self.monosaccharide_tokens)} tokens") def rule1_sia_beta_linkage(self, seq_data: Dict) -> Optional[Dict]: """Rule 1: Sialic acid with β-linkage (sialyltransferases only produce α).""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) # Find sialic acid positions sia_positions = [i for i, tid in enumerate(token_ids) if tid in self.sialic_tokens] if sia_positions and self.beta_linkage_tokens: # Change the linkage near Sia to beta pos = random.choice(sia_positions) if pos + 1 < len(token_ids): token_ids[pos + 1] = random.choice(self.beta_linkage_tokens) elif self.sialic_tokens and self.beta_linkage_tokens and len(token_ids) > 10: # Insert Sia with beta linkage pos = random.randint(3, len(token_ids) - 3) token_ids.insert(pos, random.choice(self.sialic_tokens)) token_ids.insert(pos + 1, random.choice(self.beta_linkage_tokens)) else: return None mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule1_sia_beta_linkage' return mutated def rule2_idoa_donor(self, seq_data: Dict) -> Optional[Dict]: """Rule 2: IdoA as glycosyl donor (only formed by epimerization).""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) if len(token_ids) < 8: return None # Add IdoA at a random position as if it were a donor pos = random.randint(2, len(token_ids) - 2) if self.idoa_tokens: token_ids.insert(pos, random.choice(self.idoa_tokens)) else: # Create invalid pattern token_ids[pos] = random.randint(100, self.vocab_size - 1) mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule2_idoa_donor' return mutated def rule3_high_linkage(self, seq_data: Dict) -> Optional[Dict]: """Rule 3: Linkage positions ≥8 (except α2-8 polySia).""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) if len(token_ids) < 8: return None # Insert impossible high linkage number pos = random.randint(2, len(token_ids) - 2) # Use digit tokens for 8, 9 high_digits = [tid for t, tid in self.token_to_id.items() if t in ['8', '9']] if high_digits: token_ids.insert(pos, random.choice(high_digits)) token_ids.insert(pos + 1, random.choice(high_digits)) mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule3_high_linkage' return mutated def rule4_overbranching(self, seq_data: Dict) -> Optional[Dict]: """Rule 4: More than 4 branches per residue.""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) if len(token_ids) < 10: return None # Add multiple branch tokens branch_open = self.token_to_id.get('[BRANCH_OPEN]', 5) branch_close = self.token_to_id.get('[BRANCH_CLOSE]', 6) pos = random.randint(3, len(token_ids) - 5) for _ in range(5): # 5 branches = impossible token_ids.insert(pos, branch_open) for _ in range(5): token_ids.append(branch_close) mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule4_overbranching' return mutated def rule5_sia_extension(self, seq_data: Dict) -> Optional[Dict]: """Rule 5: Extending past sialic acid (no GT accepts Sia as acceptor).""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) # Find sialic acid sia_positions = [i for i, tid in enumerate(token_ids) if tid in self.sialic_tokens] if sia_positions: # Add sugar AFTER sialic acid (impossible) pos = random.choice(sia_positions) if self.monosaccharide_tokens and pos + 2 < len(token_ids): # Insert a monosaccharide after Sia token_ids.insert(pos + 1, random.choice(self.monosaccharide_tokens)) if self.alpha_linkage_tokens: token_ids.insert(pos + 2, random.choice(self.alpha_linkage_tokens)) else: # No Sia found - add Sia then extend it if self.sialic_tokens and self.monosaccharide_tokens and len(token_ids) > 10: pos = random.randint(5, len(token_ids) - 5) token_ids.insert(pos, random.choice(self.sialic_tokens)) token_ids.insert(pos + 1, random.choice(self.monosaccharide_tokens)) else: return None mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule5_sia_extension' return mutated def rule6_c5_linkage(self, seq_data: Dict) -> Optional[Dict]: """Rule 6: C5 ring oxygen linkage (impossible).""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) if len(token_ids) < 8: return None # Insert "5" as linkage position five_token = self.token_to_id.get('5', 24) pos = random.randint(2, len(token_ids) - 2) token_ids.insert(pos, five_token) mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule6_c5_linkage' return mutated def rule7_multi_anomeric(self, seq_data: Dict) -> Optional[Dict]: """Rule 7: Multiple anomeric bonds from same carbon.""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) if len(token_ids) < 8: return None # Add conflicting anomeric markers pos = random.randint(2, len(token_ids) - 2) if self.alpha_linkage_tokens and self.beta_linkage_tokens: token_ids.insert(pos, random.choice(self.alpha_linkage_tokens)) token_ids.insert(pos + 1, random.choice(self.beta_linkage_tokens)) mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule7_multi_anomeric' return mutated def rule8_c1_acetal(self, seq_data: Dict) -> Optional[Dict]: """Rule 8: Free reducing end conflict with glycosidic bond.""" mutated = deepcopy(seq_data) token_ids = list(mutated.get('token_ids', [])) if len(token_ids) < 8: return None # Add anomeric configuration at end pos = len(token_ids) - 2 if self.alpha_linkage_tokens: token_ids.insert(pos, random.choice(self.alpha_linkage_tokens)) mutated['token_ids'] = token_ids[:256] mutated['is_negative'] = True mutated['negative_method'] = 'rule8_c1_acetal' return mutated def generate_negatives(self, n_samples: int = 100000) -> Tuple[List[Dict], Dict]: """Generate negatives with equal distribution across all 8 rules.""" logger.info(f"Generating {n_samples} negatives with 8 Universal Rules...") rules = [ (self.rule1_sia_beta_linkage, 'rule1_sia_beta'), (self.rule2_idoa_donor, 'rule2_idoa_donor'), (self.rule3_high_linkage, 'rule3_high_linkage'), (self.rule4_overbranching, 'rule4_overbranching'), (self.rule5_sia_extension, 'rule5_sia_extension'), (self.rule6_c5_linkage, 'rule6_c5_linkage'), (self.rule7_multi_anomeric, 'rule7_multi_anomeric'), (self.rule8_c1_acetal, 'rule8_c1_acetal'), ] per_rule = n_samples // len(rules) negatives = [] stats = {name: 0 for _, name in rules} pbar = tqdm(total=n_samples, desc="Generating 8-rule negatives") for rule_fn, rule_name in rules: count = 0 attempts = 0 max_attempts = per_rule * 5 while count < per_rule and attempts < max_attempts: source = random.choice(self.sequences) neg = rule_fn(source) if neg: negatives.append(neg) stats[rule_name] += 1 count += 1 pbar.update(1) attempts += 1 pbar.close() logger.info(f"Generated {len(negatives)} negatives") logger.info(f"Distribution: {json.dumps(stats, indent=2)}") return negatives, stats def main(): print("=" * 70) print("V5.1-FINAL: Corrected 8 Rules with WURCS Patterns") print("=" * 70) logger.info("Loading positive sequences...") with open(POSITIVES_PATH, 'rb') as f: positives = pickle.load(f) logger.info(f"Loaded {len(positives)} positives") generator = UniversalBiologicalRulesV4(BPE_VOCAB_PATH, positives) negatives, stats = generator.generate_negatives(n_samples=N_SAMPLES) logger.info(f"Saving {len(negatives)} negatives...") with open(OUTPUT_PATH, 'wb') as f: pickle.dump(negatives, f) # Save stats stats_output = { "total_negatives": len(negatives), "source_positives": len(positives), "mutation_types": stats, "version": "v4_FINAL_WURCS", "fixes": [ "Sialic acid detection now uses WURCS pattern a2122h/a2212h", "All token matching updated for WURCS format" ] } stats_path = OUTPUT_PATH.replace('.pkl', '_stats.json') with open(stats_path, 'w') as f: json.dump(stats_output, f, indent=2) print() print("=" * 70) print(f"Generated {len(negatives):,} rule-based negatives") print(f"Saved to: {OUTPUT_PATH}") print("=" * 70) for rule, count in sorted(stats.items()): pct = 100 * count / len(negatives) if negatives else 0 print(f" {rule}: {count:,} ({pct:.1f}%)") if __name__ == '__main__': main()