jorwal 1.0.0__py3-none-any.whl

admet_critic.py

@@ -0,0 +1,122 @@
+"""
+Agentic BioNeMo - Agent 3: ADMET & MedChem Critic Agent
+Rigorous RDKit-based physicochemical property profiling, Lipinski/Veber rules,
+PAINS structural alerts, and ADMET toxicity heuristics.
+"""
+import logging
+from typing import List, Tuple
+from rdkit import Chem
+from rdkit.Chem import Descriptors, Crippen, rdMolDescriptors, QED
+from src.models import MoleculeCandidate, AgentMessage
+
+logger = logging.getLogger("ADMETCritic")
+
+# PAINS and reactive motif SMARTS patterns
+PAINS_PATTERNS = {
+    "Quinone": "[#6]1(=[O,S])[#6]=[#6][#6](=[O,S])[#6]=[#6]1",
+    "Rhodanine": "O=C1NC(=S)SC1",
+    "Michael Acceptor": "[#6]=[#6]-[#6](=[O,S])",
+    "Aliphatic Halide": "[CX4][Cl,Br,I]",
+    "Epoxide": "C1OC1",
+    "Catechol": "c1c([OX2H])c([OX2H])ccc1"
+}
+
+class ADMETCriticAgent:
+    def __init__(self, name: str = "ADMETCritic"):
+        self.name = name
+        self._compiled_pains = {k: Chem.MolFromSmarts(v) for k, v in PAINS_PATTERNS.items() if Chem.MolFromSmarts(v)}
+
+    def evaluate_candidates(self, candidates: List[MoleculeCandidate]) -> Tuple[List[MoleculeCandidate], AgentMessage]:
+        """Runs full ADMET screening across all generated molecule candidates."""
+        passed_count = 0
+        flagged_count = 0
+        rejected_count = 0
+        
+        for cand in candidates:
+            mol = Chem.MolFromSmiles(cand.smiles)
+            if not mol:
+                cand.admet_verdict = "REJECT"
+                cand.admet_notes.append("Invalid chemical structure / SMILES parse error.")
+                rejected_count += 1
+                continue
+                
+            try:
+                Chem.SanitizeMol(mol)
+            except Exception as e:
+                cand.admet_verdict = "REJECT"
+                cand.admet_notes.append(f"Valence sanitization failed: {e}")
+                rejected_count += 1
+                continue
+                
+            # Descriptors
+            cand.mw = round(Descriptors.MolWt(mol), 2)
+            cand.logp = round(Crippen.MolLogP(mol), 2)
+            cand.tpsa = round(rdMolDescriptors.CalcTPSA(mol), 2)
+            cand.hbd = rdMolDescriptors.CalcNumHBD(mol)
+            cand.hba = rdMolDescriptors.CalcNumHBA(mol)
+            cand.rotatable_bonds = rdMolDescriptors.CalcNumRotatableBonds(mol)
+            cand.qed = round(QED.qed(mol), 3)
+            cand.molecular_formula = rdMolDescriptors.CalcMolFormula(mol)
+            
+            # Synthetic Accessibility Heuristic (Ring complexity + sp3 ratio + rotatable bonds)
+            fsp3 = rdMolDescriptors.CalcFractionCSP3(mol)
+            ring_count = rdMolDescriptors.CalcNumRings(mol)
+            # SAScore heuristic scale: 1 (easy) to 10 (hard)
+            sascore = 2.0 + (cand.mw / 150.0) + (ring_count * 0.4) - (fsp3 * 1.5)
+            cand.sascore = round(max(1.0, min(10.0, sascore)), 2)
+            
+            # PAINS alerts
+            cand.pains_alerts = []
+            for p_name, p_mol in self._compiled_pains.items():
+                if mol.HasSubstructMatch(p_mol):
+                    cand.pains_alerts.append(p_name)
+                    
+            # Lipinski & Veber
+            lipinski_violations = 0
+            if cand.mw > 550: lipinski_violations += 1
+            if cand.logp > 5.0: lipinski_violations += 1
+            if cand.hbd > 5: lipinski_violations += 1
+            if cand.hba > 10: lipinski_violations += 1
+            cand.passes_lipinski = (lipinski_violations <= 1)
+            
+            cand.passes_veber = (cand.rotatable_bonds <= 10 and cand.tpsa <= 140.0)
+            
+            # Blood-Brain Barrier (BBB) Heuristic
+            cand.bbb_permeable = (1.5 <= cand.logp <= 3.8 and cand.tpsa < 90.0 and cand.mw < 450)
+            
+            # hERG Cardiotoxicity Alert Heuristic (High LogP + high MW + basic amines)
+            cand.herg_liability = (cand.logp > 4.2 and cand.mw > 480 and cand.hba >= 4)
+            
+            # Final Verdict
+            cand.admet_notes = []
+            if cand.pains_alerts:
+                cand.admet_verdict = "REJECT"
+                cand.admet_notes.append(f"PAINS alerts: {', '.join(cand.pains_alerts)}")
+                rejected_count += 1
+            elif not cand.passes_lipinski or not cand.passes_veber:
+                cand.admet_verdict = "FLAGGED"
+                cand.admet_notes.append(f"Lipinski/Veber boundary: MW={cand.mw}, LogP={cand.logp}, TPSA={cand.tpsa}")
+                flagged_count += 1
+            elif cand.herg_liability:
+                cand.admet_verdict = "FLAGGED"
+                cand.admet_notes.append(f"hERG potential cardiotoxicity liability (LogP={cand.logp}, MW={cand.mw})")
+                flagged_count += 1
+            else:
+                cand.admet_verdict = "PASS"
+                cand.admet_notes.append(f"Compliant: QED={cand.qed}, SAScore={cand.sascore}, TPSA={cand.tpsa}")
+                passed_count += 1
+                
+        thought_text = (
+            f"Evaluated {len(candidates)} candidates through medicinal chemistry filtration. "
+            f"Outcome: {passed_count} PASS, {flagged_count} FLAGGED, {rejected_count} REJECT. "
+            f"Screened PAINS, Lipinski Rule of 5, Veber criteria, SAScore, and hERG liability."
+        )
+        message = AgentMessage(
+            agent_name=self.name,
+            role="ADMET Critic",
+            action="ADMET_MEDCHEM_EVALUATION",
+            thought=thought_text,
+            output_summary=f"{passed_count} PASS + {flagged_count} FLAGGED leads cleared for molecular docking simulation.",
+            status="SUCCESS"
+        )
+        return candidates, message

baseline.py

@@ -0,0 +1,222 @@
+"""
+Local CPU Baseline Benchmark Runner for Protein Embedding Generation.
+
+Benchmarks Hugging Face's lightweight `facebook/esm2_t6_8M_UR50D` model on local CPU
+without GPU acceleration to evaluate time-per-residue (ms/residue), token throughput
+(tokens/sec), and memory footprint.
+"""
+
+from __future__ import annotations
+
+import time
+import logging
+import tracemalloc
+from typing import Dict, Any, Optional, Tuple
+from dataclasses import dataclass, field
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_BASELINE_MODEL = "facebook/esm2_t6_8M_UR50D"
+
+
+@dataclass
+class BaselineProfileResult:
+    """Latency, throughput, and memory profiling metrics for CPU baseline inference."""
+    sequence_length: int
+    tokenization_ms: float
+    forward_pass_ms: float
+    total_cpu_latency_ms: float
+    ms_per_residue: float
+    tokens_per_sec: float
+    peak_memory_mb: float
+    tensor_bytes: int
+    tensor_shape: Tuple[int, ...]
+    model_name: str
+    is_mock: bool = False
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+class LocalCPUBaseline:
+    """
+    Runner for local CPU ESM-2 inference benchmarking.
+    Measures baseline unaccelerated compute performance across varying protein sequence lengths.
+    """
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_BASELINE_MODEL,
+        device: str = "cpu",
+        lazy_load: bool = True,
+        mock_fallback: bool = True,
+    ):
+        self.model_name = model_name
+        self.device = torch.device(device)
+        self.mock_fallback = mock_fallback
+        self.tokenizer = None
+        self.model = None
+        self._is_loaded = False
+        self._using_mock = False
+
+        if not lazy_load:
+            self.load_model()
+
+    def load_model(self) -> None:
+        """Loads tokenizer and model weights onto the target device (CPU)."""
+        if self._is_loaded:
+            return
+
+        try:
+            logger.info("Loading baseline ESM-2 model '%s' on %s...", self.model_name, self.device)
+            from transformers import AutoTokenizer, AutoModel
+
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            self.model = AutoModel.from_pretrained(self.model_name)
+            self.model.to(self.device)
+            self.model.eval()
+            self._is_loaded = True
+            self._using_mock = False
+            logger.info("Baseline ESM-2 model successfully loaded.")
+        except Exception as e:
+            if self.mock_fallback:
+                logger.warning(
+                    "Unable to load '%s' from HuggingFace hub (%s). Falling back to CPU simulation mode.",
+                    self.model_name,
+                    str(e),
+                )
+                self._is_loaded = True
+                self._using_mock = True
+            else:
+                raise RuntimeError(f"Failed to load baseline model {self.model_name}: {e}") from e
+
+    def benchmark_sequence(
+        self,
+        sequence: str,
+        warmup_runs: int = 1,
+        benchmark_runs: int = 3,
+    ) -> BaselineProfileResult:
+        """
+        Runs protein sequence embedding inference on local CPU and records execution metrics.
+
+        Args:
+            sequence: Validated amino acid sequence string.
+            warmup_runs: Initial untimed runs to warm CPU caches and instruction paths.
+            benchmark_runs: Timed repeated runs to average out timing jitter.
+
+        Returns:
+            BaselineProfileResult with detailed latency, throughput, and memory stats.
+        """
+        clean_seq = sequence.strip().upper()
+        seq_len = len(clean_seq)
+
+        if not self._is_loaded:
+            self.load_model()
+
+        if self._using_mock:
+            return self._benchmark_mock(clean_seq, seq_len)
+
+        # Warmup CPU caches
+        inputs = self.tokenizer(clean_seq, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            for _ in range(warmup_runs):
+                _ = self.model(**inputs)
+
+        # Memory profiling with tracemalloc
+        tracemalloc.start()
+        start_mem_peak = tracemalloc.get_traced_memory()[1]
+
+        # Tokenization timing
+        t_tok_start = time.perf_counter()
+        tokenized_inputs = self.tokenizer(clean_seq, return_tensors="pt")
+        tokenized_inputs = {k: v.to(self.device) for k, v in tokenized_inputs.items()}
+        t_tok_end = time.perf_counter()
+        tokenization_ms = (t_tok_end - t_tok_start) * 1000.0
+
+        # Timed forward passes
+        forward_times = []
+        last_hidden_state = None
+
+        with torch.no_grad():
+            for _ in range(max(1, benchmark_runs)):
+                t_fwd_start = time.perf_counter()
+                outputs = self.model(**tokenized_inputs)
+                t_fwd_end = time.perf_counter()
+                forward_times.append((t_fwd_end - t_fwd_start) * 1000.0)
+                last_hidden_state = outputs.last_hidden_state
+
+        current_mem, peak_mem = tracemalloc.get_traced_memory()
+        tracemalloc.stop()
+
+        peak_memory_mb = max(0.01, (peak_mem - start_mem_peak) / (1024.0 * 1024.0))
+
+        # Average forward pass duration
+        avg_forward_ms = sum(forward_times) / len(forward_times)
+        total_latency_ms = tokenization_ms + avg_forward_ms
+
+        # Throughput calculations
+        ms_per_residue = total_latency_ms / max(1, seq_len)
+        tokens_per_sec = (seq_len / (total_latency_ms / 1000.0)) if total_latency_ms > 0 else 0.0
+
+        # Tensor memory footprint
+        tensor_bytes = 0
+        tensor_shape = (1, seq_len, 320)
+        if last_hidden_state is not None:
+            tensor_bytes = last_hidden_state.element_size() * last_hidden_state.nelement()
+            tensor_shape = tuple(last_hidden_state.shape)
+
+        return BaselineProfileResult(
+            sequence_length=seq_len,
+            tokenization_ms=round(tokenization_ms, 3),
+            forward_pass_ms=round(avg_forward_ms, 3),
+            total_cpu_latency_ms=round(total_latency_ms, 3),
+            ms_per_residue=round(ms_per_residue, 4),
+            tokens_per_sec=round(tokens_per_sec, 2),
+            peak_memory_mb=round(peak_memory_mb, 2),
+            tensor_bytes=tensor_bytes,
+            tensor_shape=tensor_shape,
+            model_name=self.model_name,
+            is_mock=False,
+            metadata={"benchmark_runs": benchmark_runs, "device": str(self.device)},
+        )
+
+    def _benchmark_mock(self, clean_seq: str, seq_len: int) -> BaselineProfileResult:
+        """
+        Synthesizes realistic CPU unaccelerated baseline latency metrics.
+        On an average modern x86 CPU, 8M ESM-2 runs at roughly 0.45 - 1.2 ms per residue
+        with quadratic attention growth overhead at longer sequences.
+        """
+        # Baseline CPU performance characteristics
+        tok_ms = 0.5 + (seq_len * 0.005)
+        # O(N) feedforward + O(N^2) unaccelerated attention component
+        fwd_ms = (seq_len * 0.55) + (0.0004 * (seq_len**2))
+        total_ms = tok_ms + fwd_ms
+
+        ms_per_res = total_ms / seq_len
+        tokens_per_sec = seq_len / (total_ms / 1000.0)
+
+        # Emulate CPU processing delay (proportional but bounded for quick benchmark execution)
+        emulated_delay = min(0.08, total_ms / 2000.0)
+        time.sleep(emulated_delay)
+
+        hidden_dim = 320  # ESM2-t6-8M hidden dimension
+        tensor_shape = (1, seq_len + 2, hidden_dim)  # includes CLS and EOS tokens
+        tensor_bytes = (seq_len + 2) * hidden_dim * 4  # float32 = 4 bytes
+        peak_mem_mb = 12.0 + (seq_len * 0.02)
+
+        return BaselineProfileResult(
+            sequence_length=seq_len,
+            tokenization_ms=round(tok_ms, 3),
+            forward_pass_ms=round(fwd_ms, 3),
+            total_cpu_latency_ms=round(total_ms, 3),
+            ms_per_residue=round(ms_per_res, 4),
+            tokens_per_sec=round(tokens_per_sec, 2),
+            peak_memory_mb=round(peak_mem_mb, 2),
+            tensor_bytes=tensor_bytes,
+            tensor_shape=tensor_shape,
+            model_name=f"{self.model_name} (CPU Simulated)",
+            is_mock=True,
+            metadata={"device": "cpu-simulated"},
+        )

bus/message_bus.py

@@ -0,0 +1,101 @@
+"""
+Agentic BioNeMo - Asynchronous Swarm Message Bus
+Manages inter-agent communication, topic routing, debate history, and real-time streaming hooks.
+"""
+from typing import Callable, List, Dict, Any, Optional
+import time
+import logging
+from datetime import datetime
+from src.models import CouncilMessage
+
+logger = logging.getLogger("SwarmMessageBus")
+
+class SwarmMessageBus:
+    def __init__(self):
+        self.subscribers: Dict[str, List[Callable[[CouncilMessage], None]]] = {}
+        self.global_listeners: List[Callable[[CouncilMessage], None]] = []
+        self.history: List[CouncilMessage] = []
+        self.active_vetoes: List[CouncilMessage] = []
+        self.clearances: List[CouncilMessage] = []
+
+    def subscribe(self, topic: str, callback: Callable[[CouncilMessage], None]):
+        """Subscribe to a specific message intent/topic."""
+        if topic not in self.subscribers:
+            self.subscribers[topic] = []
+        self.subscribers[topic].append(callback)
+
+    def add_global_listener(self, callback: Callable[[CouncilMessage], None]):
+        """Subscribe to all messages published across the entire swarm."""
+        self.global_listeners.append(callback)
+
+    def publish(
+        self,
+        agent_id: str,
+        persona_name: str,
+        avatar: str,
+        intent: str,
+        content: str,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> CouncilMessage:
+        """Publish a structured council message to the swarm."""
+        now_str = datetime.now().strftime("%H:%M:%S")
+        msg = CouncilMessage(
+            agent_id=agent_id,
+            persona_name=persona_name,
+            avatar=avatar,
+            intent=intent,
+            content=content,
+            metadata=metadata or {},
+            timestamp=time.time(),
+            timestamp_str=now_str
+        )
+        
+        self.history.append(msg)
+        
+        if intent == "VETO":
+            self.active_vetoes.append(msg)
+        elif intent == "CLEARANCE":
+            self.clearances.append(msg)
+
+        # Notify topic-specific subscribers
+        if intent in self.subscribers:
+            for cb in self.subscribers[intent]:
+                try:
+                    cb(msg)
+                except Exception as e:
+                    logger.exception(f"Error in subscriber callback for {intent}: {e}")
+
+        # Notify global listeners (e.g. UI WebSocket / SSE stream)
+        for g_cb in self.global_listeners:
+            try:
+                g_cb(msg)
+            except Exception as e:
+                logger.exception(f"Error in global listener callback: {e}")
+
+        return msg
+
+    def get_council_summary(self) -> Dict[str, Any]:
+        """Returns consolidated council statistics."""
+        return {
+            "total_messages": len(self.history),
+            "veto_count": len(self.active_vetoes),
+            "clearance_count": len(self.clearances),
+            "last_action": self.history[-1].intent if self.history else "IDLE",
+            "recent_dialogues": [
+                {
+                    "agent_id": m.agent_id,
+                    "persona": m.persona_name,
+                    "avatar": m.avatar,
+                    "intent": m.intent,
+                    "content": m.content,
+                    "time": m.timestamp_str
+                }
+                for m in self.history[-12:]
+            ]
+        }
+
+    def clear(self):
+        """Reset the message bus for a new campaign."""
+        self.history.clear()
+        self.active_vetoes.clear()
+        self.clearances.clear()