Prerequisites
This section integrates everything from the chapter. You need the InfoNCE alignment framework from Section 28.1, the molecule-text and protein-text models from Section 28.2, and the document retrieval pipeline from Section 28.3. The assistant extends the Discovery Workbench introduced in Chapter 6, so familiarity with the Workbench's modular architecture is helpful. The Model Context Protocol (MCP) server integration follows patterns from Chapter 12.
A researcher investigating a new drug target needs to do three things simultaneously: search molecular databases for potential ligands, understand the target protein's structure and function, and find relevant literature that connects molecular properties to biological outcomes. Today, these three tasks use separate tools with separate interfaces and no shared context. The multimodal research assistant unifies all three into a single system with a shared embedding space. The researcher submits a query (a molecule, a protein sequence, or a natural language question) and receives coordinated results across all modalities. This section provides a complete recipe: system architecture, modality routing, shared index construction, reranking, answer synthesis, and integration with the Discovery Workbench. Figure 28.4.1 illustrates the overall architecture before we build each component.
1. System Architecture
The multimodal research assistant has five components, each building on code from earlier sections. Figure 28.4.1 shows how a query flows through the system, from modality detection through encoding, index search, reranking, and answer synthesis.
- Query Router: classifies incoming queries by modality (Simplified Molecular-Input Line-Entry System (SMILES), protein sequence, natural language) and dispatches them to the appropriate encoder.
- Modality Encoders: ChemBERTa (a transformer pretrained on chemical SMILES strings) for molecules, ESM-2 (a protein language model pretrained on millions of protein sequences) for proteins (as of 2025, ESM-3 from EvolutionaryScale offers improved protein representations and may be preferred for new projects), SciBERT (a BERT variant pretrained on scientific publications) for text, all with projection heads mapping to a shared 512-dimensional space (Section 28.1).
- Shared Index: a unified Facebook AI Similarity Search (FAISS) index containing embeddings from all modalities, enabling cross-modal nearest-neighbor search.
- Document Backend: the scientific document pipeline (Section 28.3) for literature retrieval and question answering.
- Answer Synthesizer: a large language model (LLM) that combines results from all modalities into a coherent, cited response.
import re
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
class QueryModality(Enum):
"""Detected modality of an incoming query."""
SMILES = "smiles"
PROTEIN = "protein"
TEXT = "text"
UNKNOWN = "unknown"
def detect_modality(query: str) -> QueryModality:
"""Classify a query by its scientific modality.
Uses pattern matching to distinguish SMILES strings,
protein sequences, and natural language text.
Rules:
- SMILES: contains rings (digits), bonds (=, #),
branches, and typical atoms (C, N, O, S, c, n).
- Protein: uppercase letters only from the 20 standard
amino acid single-letter codes, length >= 10.
- Text: everything else.
"""
query = query.strip()
# SMILES detection: chemical notation patterns
smiles_chars = set("CNOSPFIBrcnospfi()[]=#@+-./%0123456789")
if (len(query) >= 3
and all(c in smiles_chars for c in query)
and any(c in query for c in "()=#")):
return QueryModality.SMILES
# Protein sequence detection: amino acid alphabet
amino_acids = set("ACDEFGHIKLMNPQRSTVWY")
if (len(query) >= 10
and all(c in amino_acids for c in query.upper())
and query == query.upper()):
return QueryModality.PROTEIN
# Default: natural language text
return QueryModality.TEXT
# Test the router
assert detect_modality("CC(=O)Oc1ccccc1C(=O)O") == QueryModality.SMILES
assert detect_modality("MKTLLILAVLCLGFAQHPETLCG") == QueryModality.PROTEIN
assert detect_modality("selective EGFR inhibitor") == QueryModality.TEXT
2. The Shared Embedding Index
The core data structure is a shared FAISS index storing embeddings from all modalities. Each entry carries its source modality and metadata (SMILES string, protein ID, or document chunk reference). When a query arrives, the appropriate modality encoder encodes and projects it into the shared space. Nearest-neighbor search then returns results from any modality, ranked by cosine similarity.
A shared embedding space is a single vector space where molecules, proteins, and text documents all have numeric representations (embeddings) on the same coordinate axes. Distances between any two items are meaningful regardless of their original modality. This property enables cross-modal retrieval: without it, a text query about "kinase inhibitors" could not be compared against a molecule's SMILES representation, since the two originate from entirely different data formats. Training separate encoders with a contrastive loss (InfoNCE, where InfoNCE is a contrastive objective that maximizes the mutual information between paired items while treating all other items in the batch as negatives) creates this shared geometry. The loss pulls paired items (a molecule and its textual description, for example) together and pushes unrelated items apart. The result is a space where semantic similarity translates into vector proximity. Use a shared embedding space whenever you need to search or compare across data types in a single operation. If all your queries and results stay within one modality (text-to-text search, for instance), a standard single-modality encoder is simpler and sufficient.
Mental Model
Think of a shared embedding space like the Dewey Decimal System in a library that stocks books, vinyl records, and film reels. Each item arrives in a completely different physical format, but the classification system assigns every item a numeric code based on its subject matter, not its format. A book about jazz, a jazz album, and a jazz documentary all end up shelved near each other because their codes are close. When you walk to the 780s (music), you find all three formats side by side. The contrastive training process is analogous to the librarians who decided that subject, not format, determines the shelf number. Without that deliberate mapping, the library would have separate floors for books, records, and films, and you would need three separate trips to research a single topic.
import numpy as np
import faiss
from dataclasses import dataclass, field
@dataclass
class IndexEntry:
"""An entry in the shared multimodal index."""
modality: str # "molecule", "protein", "document"
identifier: str # SMILES, UniProt ID, or chunk reference
display_text: str # human-readable summary
metadata: dict = field(default_factory=dict)
class SharedMultimodalIndex:
"""FAISS-backed index for cross-modal retrieval.
Stores embeddings from molecules, proteins, and documents
in a single L2 index with modality metadata.
"""
def __init__(self, dimension: int = 512):
self.dimension = dimension
self.index = faiss.IndexFlatIP(dimension) # inner product
self.entries: list[IndexEntry] = []
def add_embeddings(self, embeddings: np.ndarray,
entries: list[IndexEntry]) -> None:
"""Add a batch of embeddings with metadata.
Parameters
----------
embeddings : np.ndarray
L2-normalized vectors, shape (N, dimension).
entries : list[IndexEntry]
Metadata for each embedding.
"""
assert len(embeddings) == len(entries)
assert embeddings.shape[1] == self.dimension
# FAISS requires float32
emb32 = embeddings.astype(np.float32)
faiss.normalize_L2(emb32) # ensure normalization
self.index.add(emb32)
self.entries.extend(entries)
def search(self, query_embedding: np.ndarray,
top_k: int = 20,
modality_filter: str = None) -> list[dict]:
"""Search the index for nearest neighbors.
Parameters
----------
query_embedding : np.ndarray
L2-normalized query vector, shape (1, dimension).
top_k : int
Number of results to return.
modality_filter : str, optional
Restrict results to a specific modality.
"""
q = query_embedding.astype(np.float32).reshape(1, -1)
faiss.normalize_L2(q)
# Search more than top_k to allow for filtering
search_k = top_k * 3 if modality_filter else top_k
scores, indices = self.index.search(q, search_k)
results = []
for score, idx in zip(scores[0], indices[0]):
if idx < 0: # FAISS returns -1 for empty slots
continue
entry = self.entries[idx]
if modality_filter and entry.modality != modality_filter:
continue
results.append({
"score": float(score),
"modality": entry.modality,
"identifier": entry.identifier,
"display_text": entry.display_text,
"metadata": entry.metadata,
})
if len(results) >= top_k:
break
return results
@property
def size(self) -> int:
return self.index.ntotal
def modality_counts(self) -> dict:
"""Count entries by modality."""
counts = {}
for entry in self.entries:
counts[entry.modality] = counts.get(entry.modality, 0) + 1
return counts
modality_filter parameter restricts results to a single modality while still using the unified index.Common Misconception
A frequent mistake is assuming that cosine similarity scores are directly comparable across modalities: that a molecule result at 0.82 is more relevant than a document result at 0.78. In practice, each modality's embeddings occupy a different region of the shared space with distinct density and spread (the "modality gap"), so raw scores from different modalities are on different scales. This is exactly why the reranking stage in Section 5 applies z-score normalization within each modality before comparing results; without that normalization, the final ranking can be dominated by whichever modality happens to produce numerically higher raw scores.
A single shared index works because the alignment training (Section 28.1) ensures that semantically related items from different modalities occupy the same region of the embedding space. A text query for "selective Epidermal Growth Factor Receptor (EGFR) inhibitor" returns nearby molecule embeddings (actual EGFR inhibitor compounds), protein embeddings (the EGFR kinase domain), and document embeddings (papers about EGFR-targeted therapy). This cross-modal retrieval emerges from the alignment, not from any special index structure. The FAISS index is modality-agnostic; it operates on raw vectors. The semantic structure comes entirely from the encoders and their shared projection space.
3. The Multimodal Research Assistant
Without this orchestration layer, researchers must manually copy results between separate molecular, protein, and literature search tools, a workflow that routinely causes missed connections between promising drug candidates and their target proteins. The assistant described next eliminates that fragmentation.
With a shared index that can store and retrieve embeddings from any modality, the remaining challenge is orchestration: accepting a user query, routing it to the correct encoder, searching the index, and assembling a coherent response from the mixed results.
The MultimodalResearchAssistant class below implements this orchestration as the core of the Discovery Workbench's multimodal module. In short: one query, one embedding, all modalities searched at once.
from dataclasses import dataclass, field
from typing import Optional
import json
@dataclass
class AssistantResponse:
"""Structured response from the multimodal research assistant."""
query: str
detected_modality: str
molecular_results: list[dict] = field(default_factory=list)
protein_results: list[dict] = field(default_factory=list)
literature_results: list[dict] = field(default_factory=list)
synthesis: str = ""
confidence: float = 0.0
class MultimodalResearchAssistant:
"""A multimodal research assistant for scientific discovery.
Accepts molecular queries (SMILES), protein sequences, or
natural language questions and returns coordinated results
across molecular databases, protein annotations, and
scientific literature.
Integrates with the Discovery Workbench as a pluggable module.
"""
def __init__(
self,
shared_index: SharedMultimodalIndex,
mol_encoder, # trained molecule encoder + projection
protein_encoder, # trained protein encoder + projection
text_encoder, # trained text encoder + projection
document_pipeline, # ScientificDocumentPipeline
mol_tokenizer,
protein_tokenizer,
text_tokenizer,
llm=None, # optional LLM for synthesis
):
self.index = shared_index
self.mol_encoder = mol_encoder
self.protein_encoder = protein_encoder
self.text_encoder = text_encoder
self.doc_pipeline = document_pipeline
self.mol_tok = mol_tokenizer
self.prot_tok = protein_tokenizer
self.text_tok = text_tokenizer
self.llm = llm
def _encode_query(self, query: str,
modality: QueryModality) -> np.ndarray:
"""Encode a query using the appropriate modality encoder."""
import torch
import torch.nn.functional as F
if modality == QueryModality.SMILES:
inputs = self.mol_tok(
query, return_tensors="pt", max_length=128,
truncation=True, padding=True,
)
with torch.no_grad():
z = self.mol_encoder(**inputs)
elif modality == QueryModality.PROTEIN:
inputs = self.prot_tok(
query, return_tensors="pt", max_length=1024,
truncation=True, padding=True,
)
with torch.no_grad():
z = self.protein_encoder(**inputs)
else: # TEXT
inputs = self.text_tok(
query, return_tensors="pt", max_length=256,
truncation=True, padding=True,
)
with torch.no_grad():
z = self.text_encoder(**inputs)
z = F.normalize(z, dim=-1)
return z.cpu().numpy()
def query(self, query: str,
top_k_per_modality: int = 5) -> AssistantResponse:
"""Process a multimodal query and return coordinated results.
Parameters
----------
query : str
A SMILES string, protein sequence, or natural language.
top_k_per_modality : int
Number of results per modality.
"""
modality = detect_modality(query)
z = self._encode_query(query, modality)
# Search the shared index across all modalities
all_results = self.index.search(z, top_k=top_k_per_modality * 3)
# Partition results by modality
mol_results = [r for r in all_results
if r["modality"] == "molecule"][:top_k_per_modality]
prot_results = [r for r in all_results
if r["modality"] == "protein"][:top_k_per_modality]
doc_results = [r for r in all_results
if r["modality"] == "document"][:top_k_per_modality]
# Augment with document pipeline for text queries
if modality == QueryModality.TEXT and self.doc_pipeline:
lit_results = self.doc_pipeline.query(
query, top_k=top_k_per_modality
)
# Merge with cross-modal document results
doc_results = self._merge_results(doc_results, lit_results)
response = AssistantResponse(
query=query,
detected_modality=modality.value,
molecular_results=mol_results,
protein_results=prot_results,
literature_results=doc_results,
)
# Synthesize if LLM is available
if self.llm:
response.synthesis = self._synthesize(response)
response.confidence = self._estimate_confidence(response)
return response
def _merge_results(self, cross_modal: list[dict],
direct: list[dict]) -> list[dict]:
"""Merge cross-modal and direct retrieval results.
Deduplicates and re-ranks by combining scores.
"""
seen = set()
merged = []
for r in cross_modal + direct:
key = r.get("identifier", r.get("content", ""))[:100]
if key not in seen:
seen.add(key)
merged.append(r)
# Sort by score (descending)
merged.sort(key=lambda x: x.get("score", 0), reverse=True)
return merged
def _synthesize(self, response: AssistantResponse) -> str:
"""Generate a synthesis from cross-modal results."""
prompt = self._build_synthesis_prompt(response)
# In production, call the LLM here
# return self.llm.generate(prompt)
return f"[Synthesis based on {len(response.molecular_results)} molecules, " \
f"{len(response.protein_results)} proteins, " \
f"{len(response.literature_results)} literature sources]"
def _build_synthesis_prompt(self,
response: AssistantResponse) -> str:
"""Build the LLM prompt for answer synthesis."""
sections = [f"Query: {response.query}\n"]
if response.molecular_results:
sections.append("Relevant Molecules:")
for r in response.molecular_results[:3]:
sections.append(
f" - {r['identifier']} (similarity: {r['score']:.3f})"
f"\n {r['display_text']}"
)
if response.protein_results:
sections.append("\nRelevant Proteins:")
for r in response.protein_results[:3]:
sections.append(
f" - {r['identifier']} (similarity: {r['score']:.3f})"
f"\n {r['display_text']}"
)
if response.literature_results:
sections.append("\nRelevant Literature:")
for r in response.literature_results[:3]:
sections.append(
f" - {r.get('display_text', r.get('content', ''))[:200]}"
)
sections.append(
"\nSynthesize the above evidence into a concise answer. "
"Cite specific molecules, proteins, and papers. "
"Identify gaps where additional evidence is needed."
)
return "\n".join(sections)
def _estimate_confidence(self,
response: AssistantResponse) -> float:
"""Estimate response confidence from result quality.
Higher confidence when multiple modalities agree and
top results have high similarity scores.
"""
scores = []
for results in [response.molecular_results,
response.protein_results,
response.literature_results]:
if results:
scores.append(results[0].get("score", 0))
if not scores:
return 0.0
# Confidence: average top score * coverage factor
avg_score = sum(scores) / len(scores)
coverage = len(scores) / 3 # fraction of modalities with results
return min(1.0, avg_score * coverage)
Note that the assistant's _merge_results method sorts by raw similarity scores, which are not yet calibrated across modalities. Section 5 introduces a reranking stage with z-score normalization that should be applied to the combined output before presenting results to the user; in a production system, you would call the reranker between the index search and the answer synthesis steps.
Checkpoint
So far: the assistant accepts a query, routes it to the correct modality encoder, searches a shared FAISS index for nearest neighbors across molecules, proteins, and documents, merges results from cross-modal and direct retrieval, and optionally synthesizes a cited answer via an LLM.
A researcher at a biotech startup queries the assistant with the SMILES string for
sorafenib (a multi-kinase inhibitor): CNC(=O)c1cc(Oc2ccc(NC(=O)Nc3ccc(Cl)c(C(F)(F)F)c3)cc2)ccn1.
The assistant detects the SMILES modality and returns:
Molecular results: 5 structurally similar compounds including regorafenib
and lenvatinib, both FDA-approved kinase inhibitors.
Protein results: VEGFR2, RAF1, and BRAF kinases appear as the nearest protein
embeddings, consistent with sorafenib's known target profile.
Literature results: Papers on sorafenib resistance mechanisms, combination
therapy with immune checkpoint inhibitors, and a 2024 study on sorafenib analogs with
improved selectivity.
The synthesis identifies a gap: no literature discusses the combination of sorafenib
analogs with the specific BRAF mutants found in the protein results, suggesting a
testable hypothesis for combination therapy in BRAF-mutant cancers. This hypothesis
feeds directly into the hypothesis generation
pipeline in Chapter 39.
Exercise 28.4.1
The detect_modality function checks for SMILES patterns before checking for
protein sequences. Suppose you swap the order so that protein detection runs first.
Which of the following queries would be misclassified, and why:
(a) CNOS, (b) CC(=O)O, (c) ACDEFGHIKLMNPQRSTVWY,
(d) selective kinase inhibitor?
Hint
Consider the protein detection rule: all characters must be uppercase letters from the 20 standard amino acid codes, and the string must be at least 10 characters long. Ask yourself which of the four queries satisfies those conditions. Then check whether any of those same queries also satisfy the SMILES conditions. The answer hinges on a query that is valid under both rules.
4. Populating the Shared Index
The assistant is only as good as its index. Populating the shared index requires encoding large molecular databases, protein databases, and document corpora. The following code shows the batch encoding pipeline for molecules (sourced from PubChem, the public repository of chemical structures and their biological activities maintained by the National Center for Biotechnology Information) and proteins (sourced from UniProt, the comprehensive database of protein sequence and functional information).
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
def populate_molecular_index(
smiles_list: list[str],
descriptions: list[str],
mol_encoder,
mol_tokenizer,
index: SharedMultimodalIndex,
batch_size: int = 64,
) -> int:
"""Encode and index a molecular database.
Parameters
----------
smiles_list : list[str]
SMILES strings for all molecules.
descriptions : list[str]
Brief text descriptions for display.
mol_encoder : nn.Module
Trained molecule encoder with projection head.
mol_tokenizer : tokenizer
Tokenizer for the molecule encoder.
index : SharedMultimodalIndex
The shared index to populate.
batch_size : int
Encoding batch size (adjust for GPU memory).
Returns
-------
int
Number of molecules added to the index.
"""
all_embeddings = []
entries = []
for i in tqdm(range(0, len(smiles_list), batch_size),
desc="Encoding molecules"):
batch_smiles = smiles_list[i:i + batch_size]
batch_descs = descriptions[i:i + batch_size]
inputs = mol_tokenizer(
batch_smiles, return_tensors="pt", padding=True,
truncation=True, max_length=128,
)
with torch.no_grad():
z = mol_encoder(**inputs)
z = torch.nn.functional.normalize(z, dim=-1)
all_embeddings.append(z.cpu().numpy())
for smi, desc in zip(batch_smiles, batch_descs):
entries.append(IndexEntry(
modality="molecule",
identifier=smi,
display_text=desc,
))
embeddings = np.concatenate(all_embeddings, axis=0)
index.add_embeddings(embeddings, entries)
return len(entries)
def populate_protein_index(
sequences: list[str],
protein_ids: list[str],
annotations: list[str],
protein_encoder,
protein_tokenizer,
index: SharedMultimodalIndex,
batch_size: int = 16,
) -> int:
"""Encode and index a protein database.
Similar to molecular indexing but with larger sequences
requiring smaller batch sizes for memory management.
"""
all_embeddings = []
entries = []
for i in tqdm(range(0, len(sequences), batch_size),
desc="Encoding proteins"):
batch_seqs = sequences[i:i + batch_size]
batch_ids = protein_ids[i:i + batch_size]
batch_annots = annotations[i:i + batch_size]
inputs = protein_tokenizer(
batch_seqs, return_tensors="pt", padding=True,
truncation=True, max_length=1024,
)
with torch.no_grad():
z = protein_encoder(**inputs)
z = torch.nn.functional.normalize(z, dim=-1)
all_embeddings.append(z.cpu().numpy())
for pid, annot in zip(batch_ids, batch_annots):
entries.append(IndexEntry(
modality="protein",
identifier=pid,
display_text=annot,
))
embeddings = np.concatenate(all_embeddings, axis=0)
index.add_embeddings(embeddings, entries)
return len(entries)
5. Reranking Cross-Modal Results
Populating the index gives the assistant something to search, but the raw similarity scores it returns carry a subtle problem: they are not calibrated across modalities.
Raw cross-modal retrieval scores are not directly comparable across modalities. A molecule at cosine similarity 0.75 and a document at 0.80 may not reflect true relative relevance because the embedding distributions differ by modality (the "modality gap," a systematic offset between the regions of the shared space occupied by different modalities, as discussed in Section 28.1). A reranking stage normalizes scores and applies Maximal Marginal Relevance (MMR), a selection strategy that balances relevance to the query against diversity among the selected results, penalizing candidates that are too similar to items already chosen. The normalization uses z-score transformation, where each raw score is converted to the number of standard deviations it lies above or below the mean score for its modality, making scores from different distributions comparable.
from dataclasses import dataclass
@dataclass
class RerankerConfig:
"""Configuration for cross-modal result reranking."""
mol_weight: float = 1.0 # modality importance weight
protein_weight: float = 1.0
document_weight: float = 1.0
diversity_penalty: float = 0.1 # penalize redundant results
def rerank_results(
mol_results: list[dict],
protein_results: list[dict],
doc_results: list[dict],
config: RerankerConfig = RerankerConfig(),
top_k: int = 10,
) -> list[dict]:
"""Rerank cross-modal results for unified presentation.
Normalizes scores within each modality (z-score),
applies modality weights, and penalizes redundancy
using maximal marginal relevance (MMR).
Parameters
----------
mol_results, protein_results, doc_results : list[dict]
Results from each modality with "score" fields.
config : RerankerConfig
Modality weights and diversity penalty.
top_k : int
Total number of results to return.
"""
# Z-score normalization within each modality
def normalize_scores(results: list[dict]) -> list[dict]:
if not results:
return results
scores = [r["score"] for r in results]
mean_s = sum(scores) / len(scores)
std_s = (sum((s - mean_s)**2 for s in scores)
/ len(scores)) ** 0.5
if std_s < 1e-6:
std_s = 1.0
for r in results:
r["normalized_score"] = (r["score"] - mean_s) / std_s
return results
mol_results = normalize_scores(mol_results)
protein_results = normalize_scores(protein_results)
doc_results = normalize_scores(doc_results)
# Apply modality weights
weights = {
"molecule": config.mol_weight,
"protein": config.protein_weight,
"document": config.document_weight,
}
all_results = []
for results, modality in [
(mol_results, "molecule"),
(protein_results, "protein"),
(doc_results, "document"),
]:
for r in results:
r["weighted_score"] = (
r.get("normalized_score", r["score"])
* weights[modality]
)
r["modality"] = modality
all_results.append(r)
# Sort by weighted score
all_results.sort(key=lambda x: x["weighted_score"], reverse=True)
# MMR-style diversity: penalize results similar to already-selected ones
selected = []
for candidate in all_results:
if len(selected) >= top_k:
break
# Penalize if same modality and similar identifier
penalty = 0.0
for s in selected:
if s["modality"] == candidate["modality"]:
penalty += config.diversity_penalty
candidate["final_score"] = (
candidate["weighted_score"] - penalty
)
selected.append(candidate)
selected.sort(key=lambda x: x["final_score"], reverse=True)
return selected
Step-Through: Cross-Modal Reranking
Trace through the reranking algorithm with three results, one per modality. Raw scores: molecule A = 0.82, protein B = 0.61, document C = 0.74. Z-score normalization (within each single-element modality group, std defaults to 1.0): normalized molecule A = (0.82 − 0.82)/1.0 = 0.0, protein B = 0.0, document C = 0.0. All normalized scores are identical, so modality weights (all 1.0) leave them tied at 0.0. MMR diversity pass (penalty = 0.1): the first selected result gets no penalty (final = 0.0). The second candidate shares no modality with the first, so penalty = 0.0 and final = 0.0. The third candidate shares no modality with either selected result, so penalty = 0.0 again. All three survive with final_score = 0.0. Now change the scenario: add a second molecule A2 with raw score 0.79. Within the molecule group, mean = (0.82 + 0.79)/2 = 0.805, std = 0.015. Normalized: A = (0.82 − 0.805)/0.015 = 1.0, A2 = (0.79 − 0.805)/0.015 = −1.0. After the MMR pass selects A first, A2 incurs a 0.1 same-modality penalty, dropping to −1.1, which pushes it below protein B and document C in the final ranking.
6. Integration with the Discovery Workbench
With query routing, shared indexing, and cross-modal reranking in place, the assistant is functionally complete; what remains is exposing it as a composable service inside the broader Discovery Workbench.
The multimodal research assistant integrates with the Discovery Workbench as a module that can be invoked from the Workbench's agent framework. Following the MCP server patterns from Chapter 12, the assistant exposes its capabilities as tool calls that agents can compose into complex discovery workflows (the research agent patterns in Chapter 40 generalize this approach).
from typing import Any
class MultimodalWorkbenchModule:
"""Discovery Workbench module for multimodal scientific queries.
Exposes the multimodal research assistant as a set of
tool functions callable from Workbench agents and pipelines.
"""
def __init__(self, assistant: MultimodalResearchAssistant):
self.assistant = assistant
self.tools = {
"multimodal_search": self.multimodal_search,
"molecule_search": self.molecule_search,
"protein_search": self.protein_search,
"literature_search": self.literature_search,
"cross_modal_analysis": self.cross_modal_analysis,
}
def multimodal_search(self, query: str,
top_k: int = 5) -> dict:
"""Search across all modalities with a single query.
Accepts SMILES, protein sequences, or natural language.
Returns results from molecules, proteins, and literature.
"""
response = self.assistant.query(query, top_k)
return {
"query": response.query,
"modality": response.detected_modality,
"molecules": response.molecular_results,
"proteins": response.protein_results,
"literature": response.literature_results,
"synthesis": response.synthesis,
"confidence": response.confidence,
}
def molecule_search(self, query: str,
top_k: int = 10) -> list[dict]:
"""Search only molecular results."""
z = self.assistant._encode_query(
query, detect_modality(query)
)
return self.assistant.index.search(
z, top_k=top_k, modality_filter="molecule"
)
def protein_search(self, query: str,
top_k: int = 10) -> list[dict]:
"""Search only protein results."""
z = self.assistant._encode_query(
query, detect_modality(query)
)
return self.assistant.index.search(
z, top_k=top_k, modality_filter="protein"
)
def literature_search(self, query: str,
top_k: int = 10) -> list[dict]:
"""Search the literature corpus."""
if self.assistant.doc_pipeline:
return self.assistant.doc_pipeline.query(query, top_k)
return []
def cross_modal_analysis(self, smiles: str,
protein_id: str) -> dict:
"""Analyze the relationship between a molecule and protein.
Returns similarity in the shared space, relevant literature,
and a synthesis of what is known about the interaction.
"""
z_mol = self.assistant._encode_query(
smiles, QueryModality.SMILES
)
# Encode the protein identifier as text (not as a raw sequence),
# since protein_id is a name like "EGFR_HUMAN", not a full
# amino acid sequence. For sequence-level analysis, pass the
# actual sequence and use QueryModality.PROTEIN instead.
z_prot = self.assistant._encode_query(
protein_id, QueryModality.TEXT
)
# Cross-modal similarity
similarity = float(
np.dot(z_mol.squeeze(), z_prot.squeeze())
)
# Find literature about this molecule-protein pair
combined_query = (
f"interaction between {smiles} and {protein_id}"
)
lit_results = self.literature_search(combined_query, top_k=5)
return {
"molecule": smiles,
"protein": protein_id,
"cross_modal_similarity": similarity,
"supporting_literature": lit_results,
"interaction_likely": similarity > 0.5, # heuristic threshold; tune on validated interaction data
}
def get_tool_descriptions(self) -> list[dict]:
"""Return MCP-compatible tool descriptions.
These descriptions are used by the Discovery Workbench
agent framework for tool selection and composition.
"""
return [
{
"name": "multimodal_search",
"description": (
"Search across molecules, proteins, and "
"literature with a single query. Accepts "
"SMILES strings, protein sequences, or "
"natural language."
),
"parameters": {
"query": {"type": "string", "required": True},
"top_k": {"type": "integer", "default": 5},
},
},
{
"name": "cross_modal_analysis",
"description": (
"Analyze the predicted interaction between "
"a molecule (SMILES) and a protein. Returns "
"cross-modal similarity and supporting literature."
),
"parameters": {
"smiles": {"type": "string", "required": True},
"protein_id": {"type": "string", "required": True},
},
},
]
get_tool_descriptions method returns JSON schemas that the Workbench agent framework uses for automatic tool discovery and composition.The assistant built in this section trains separate encoders per modality and aligns them post hoc. A newer line of work trains a single foundation model that natively handles multiple biomedical modalities. BioMedGPT (Luo et al., 2023) jointly pretrains on molecules, proteins, and biomedical text using a unified transformer architecture, reporting strong results on cross-modal retrieval benchmarks without requiring separate projection heads. More recently, BioMedParse (Zhao et al., 2024, published in Nature Methods) extended unified multimodal modeling to biomedical image segmentation across nine imaging modalities, demonstrating that a single model can parse radiology, pathology, and microscopy images with a shared representation. These unified architectures largely eliminate the modality gap problem, since all modalities are encoded by the same transformer from the start, though they require substantially more pretraining data and compute than the modular approach presented here.
Real-World Application: Pharmaceutical Target Identification
Recursion Pharmaceuticals reportedly operates a multimodal discovery platform that jointly queries cellular microscopy images, molecular structures, and genomic data through a shared embedding space (their "Recursion OS"). When a researcher submits a disease phenotype, the system is designed to retrieve candidate compounds, gene targets, and published evidence in a single pass, enabling hypothesis generation that would otherwise require separate searches across disconnected databases. The architecture resembles the shared-index pattern in this section, scaled to billions of cellular images and millions of compounds.
The from-scratch Workbench module above is approximately 180 lines. LangChain's agent framework provides tool management, memory, and orchestration that reduces the integration code to about 30 lines:
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.tools import StructuredTool
from langchain_openai import ChatOpenAI
# Wrap assistant methods as LangChain tools
tools = [
StructuredTool.from_function(
func=workbench_module.multimodal_search,
name="multimodal_search",
description="Search molecules, proteins, literature",
),
StructuredTool.from_function(
func=workbench_module.cross_modal_analysis,
name="cross_modal_analysis",
description="Analyze molecule-protein interactions",
),
]
llm = ChatOpenAI(model="gpt-4o")
agent = create_tool_calling_agent(llm, tools, prompt_template)
executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
# The agent decides which tools to call and in what order
result = executor.invoke({
"input": "Find EGFR inhibitors and check their selectivity"
})
StructuredTool.from_function, enabling automatic tool selection, argument parsing, and multi-step execution via AgentExecutor.LangChain handles tool selection, argument parsing, memory management, and multi-step execution internally. Line count reduction: approximately 6x, from 180 lines of hand-rolled orchestration down to 30. Use the from-scratch approach when you need fine-grained control over tool selection logic or custom reranking strategies; use LangChain for rapid prototyping and standard agent patterns. As of 2025, the LangChain project recommends LangGraph over AgentExecutor for new agent applications; LangGraph provides a graph-based orchestration model with stronger support for multi-step reasoning, explicit state management, and composable tool pipelines.
During testing, the modality router classified the query "CANCER" as a protein sequence (it is a valid uppercase amino acid string: C, A, N, C, E, R). The fix: require protein sequences to be at least 10 characters long and not match common English words. This edge case illustrates a general principle of multimodal systems: the boundaries between modalities are not always clear, and the router must be robust to ambiguous inputs. The word "MAGIC" is also a valid amino acid sequence. So is "DELIVER." Chemistry and biology have a way of overlapping with the English language.
7. End-to-End Example: A Complete Discovery Session
The following walkthrough demonstrates a full discovery session with the multimodal research assistant, from initial query through hypothesis generation.
def discovery_session_example():
"""Demonstrate a complete multimodal discovery session.
Scenario: a researcher investigating new treatments for
non-small cell lung cancer (NSCLC) with EGFR mutations.
"""
# Step 1: Natural language query
print("=== Step 1: Initial Literature Search ===")
lit_results = workbench.literature_search(
"EGFR inhibitor resistance mechanisms in NSCLC",
top_k=5,
)
for r in lit_results:
print(f" [{r.get('type', 'text')}] {r.get('content', '')[:100]}...")
# Step 2: Molecular search for known EGFR inhibitors
print("\n=== Step 2: Molecular Search ===")
mol_results = workbench.molecule_search(
"third generation EGFR tyrosine kinase inhibitor",
top_k=5,
)
for r in mol_results:
print(f" {r['identifier']} (sim={r['score']:.3f})")
print(f" {r['display_text']}")
# Step 3: Cross-modal analysis of osimertinib with EGFR
print("\n=== Step 3: Cross-Modal Analysis ===")
osimertinib_smiles = (
"COc1cc2ncnc(Nc3ccc(F)c(Cl)c3)c2cc1NC(=O)C=C"
)
analysis = workbench.cross_modal_analysis(
smiles=osimertinib_smiles,
protein_id="EGFR_HUMAN",
)
print(f" Similarity: {analysis['cross_modal_similarity']:.3f}")
print(f" Interaction likely: {analysis['interaction_likely']}")
print(f" Supporting papers: {len(analysis['supporting_literature'])}")
# Step 4: Multimodal search for resistance-overcoming compounds
print("\n=== Step 4: Multimodal Discovery Query ===")
discovery = workbench.multimodal_search(
"EGFR inhibitor that overcomes C797S resistance mutation",
top_k=3,
)
print(f" Molecules found: {len(discovery['molecules'])}")
print(f" Proteins found: {len(discovery['proteins'])}")
print(f" Literature found: {len(discovery['literature'])}")
print(f" Confidence: {discovery['confidence']:.2f}")
print(f"\n Synthesis: {discovery['synthesis']}")
# Step 5: Generate hypothesis
print("\n=== Step 5: Hypothesis ===")
hypothesis = (
f"Based on {len(discovery['molecules'])} molecular hits, "
f"{len(discovery['proteins'])} protein targets, and "
f"{len(discovery['literature'])} literature sources: "
f"Test the top molecular hit against the C797S mutant EGFR "
f"in a binding assay to validate the cross-modal prediction."
)
print(f" {hypothesis}")
# discovery_session_example()
Try It: Build a Mini Cross-Modal Search Engine
Build a working cross-modal retrieval system on your laptop using freely available models and data.
Step 1: Install dependencies: pip install transformers faiss-cpu numpy torch. Load the sentence-transformers/all-MiniLM-L6-v2 model as your text encoder (384 dimensions; no GPU required).
Step 2: Collect 50 molecule descriptions from PubChem by downloading the first page of compounds at https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/1-50/description/JSON. Extract the text description for each compound.
Step 3: Encode all 50 descriptions into embeddings using the sentence transformer and add them to a faiss.IndexFlatIP index. Store the compound names and CIDs alongside each embedding in a parallel list (mirroring the IndexEntry pattern from this section).
Step 4: Encode a natural language query such as "anti-inflammatory pain reliever" with the same model. Run index.search() and verify that aspirin and ibuprofen rank near the top.
Step 5: Add 20 short abstracts from PubMed (search "kinase inhibitor" and copy the first 20 results) to the same index. Re-run the query and observe that the results now include both compound descriptions and paper abstracts, demonstrating cross-modal retrieval within a single index. Measure query latency; it should be under 10 milliseconds for this index size.
Lab: Cross-Modal Retrieval Accuracy vs. Index Composition
Goal: Measure how the ratio of molecules to documents in a shared index
affects cross-modal retrieval precision for natural language queries.
Tools: Python, faiss-cpu, sentence-transformers
(the all-MiniLM-L6-v2 model), and the PubChem REST API.
Setup (15 min): Download 100 compound descriptions from PubChem
(https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/1-100/description/JSON).
Collect 100 short PubMed abstracts (search "drug discovery" via the Entrez API).
Encode all 200 items with the sentence transformer and build a FAISS IndexFlatIP index.
Experiment (15 min): Prepare five queries with known relevant compounds
(e.g., "anti-inflammatory" should retrieve ibuprofen and aspirin). Build five versions of
the index using molecule-to-document ratios of 10:90, 25:75, 50:50, 75:25, and 90:10
(sampling from your 200 items). For each ratio, run all five queries and record
precision@5 (fraction of top-5 results that are genuinely relevant).
What to observe: Does retrieval precision for molecular hits degrade when
documents dominate the index? Does the optimal ratio differ across queries? Plot
precision@5 vs. ratio and note whether the curve is symmetric or skewed toward one modality.
Exercises
- Conceptual. The modality router uses simple pattern matching to classify queries. Design a more robust routing strategy that handles ambiguous inputs (queries that could be interpreted as multiple modalities). Consider: (a) confidence scores for each modality classification, (b) running all encoders and using the embedding norms to select the best modality, (c) asking the user for clarification. What are the latency tradeoffs?
- Coding. Build a minimal multimodal research assistant using the code from this section. Use ChemBERTa (from Hugging Face) as the molecule encoder, ESM-2 as the protein encoder, and SciBERT as the text encoder. Populate the shared FAISS index with 100 molecules from PubChem, 50 protein sequences from UniProt, and 200 text chunks from 5 arXiv papers. Demonstrate a cross-modal query that retrieves results from all three modalities. Measure the total query latency and identify the bottleneck.
- Analysis. The confidence estimation in the assistant multiplies average top-score by modality coverage. Critique this approach: under what conditions does it overestimate confidence? Under what conditions does it underestimate? Propose an alternative confidence metric that accounts for (a) inter-result consistency (do the molecular and literature results agree?), (b) result diversity (are all molecular results structurally similar or diverse?), and (c) known calibration data (historical accuracy on similar queries).