28.2 Molecule-Text and Protein-Text Models

Before diving into each model, Figure 28.2 illustrates how text serves as a universal bridge modality. Separate models align molecules and proteins to text independently, and the overlapping text embedding space creates emergent cross-modal connections between molecules and proteins without requiring direct molecule-protein training pairs.

1. MolT5: Bidirectional Molecule-Text Translation

What if you could type "a selective COX-2 inhibitor with low hepatotoxicity risk" into a search bar and receive, in return, the SMILES string of a molecule that matches that description?

Without models like these, a medicinal chemist searching for "a selective COX-2 inhibitor with low hepatotoxicity risk" would have to translate that intent into a precise substructure query, manually filtering thousands of hits; any molecule described differently in the database would be invisible. Molecule-text models remove that translation step entirely.

MolT5 (Edwards et al., 2022) makes this possible by treating molecule-text translation as a sequence-to-sequence problem. Built on the T5 (Text-to-Text Transfer Transformer) architecture, it handles two tasks: molecule captioning (SMILES to text) and text-conditional molecule generation (text to SMILES). The authors pretrain the model on a large corpus of unlabeled SMILES strings and scientific text, then fine-tune it on paired molecule-text data from the CheBI-20 dataset (33,010 molecule-description pairs derived from the Chemical Entities of Biological Interest ontology). (As of 2024, BioT5 and BioT5+ extend the MolT5 paradigm by jointly pretraining on molecules, proteins, and biomedical text, achieving stronger captioning and generation scores on CheBI-20 and additional protein-related benchmarks.)

Molecule-text translation converts a machine-readable molecular representation (such as a SMILES string) into a natural language description of that molecule's structure, properties, or biological activity, and vice versa. It matters because it bridges two communities: computational chemists who work with structural formulas and bench scientists or clinicians who think in terms of drug classes, mechanisms, and side effects. An encoder-decoder transformer ingests one representation as a token sequence and autoregressively generates the other. Use molecule-text translation to search chemical databases with plain-language queries, to generate human-interpretable explanations of screening hits, or to produce candidate structures from verbal specifications. Use direct molecular property prediction (quantitative structure-activity relationship (QSAR) models) instead when you need precise numeric properties like binding affinity or solubility and already have the molecular structure in hand.

The architecture is a T5 encoder-decoder where the input is either a SMILES string (for captioning) or a text description (for generation), and the output is the other modality. No special molecular tokenizer is required; MolT5 treats SMILES characters as text tokens, though specialized tokenizers (like those splitting on chemical bonds) can improve performance. At decoding time, MolT5 uses beam search (a breadth-limited tree search that keeps the top-k most probable partial sequences at each step, trading off speed for output quality) to produce fluent captions or valid SMILES. In short: by treating SMILES and English as two dialects of the same language, a single sequence-to-sequence model can translate freely between how machines encode molecules and how scientists talk about them.

from transformers import T5ForConditionalGeneration, T5Tokenizer
import torch


def load_molt5(model_name: str = "laituan245/molt5-base") -> tuple:
    """Load MolT5 model and tokenizer from Hugging Face.

    Available sizes: molt5-small (77M), molt5-base (250M),
    molt5-large (780M).
    """
    tokenizer = T5Tokenizer.from_pretrained(model_name)
    model = T5ForConditionalGeneration.from_pretrained(model_name)
    model.eval()
    return model, tokenizer


def smiles_to_caption(smiles: str, model, tokenizer,
                      max_length: int = 256) -> str:
    """Generate a text description for a SMILES string.

    MolT5 molecule captioning: SMILES in, English out.
    """
    input_ids = tokenizer(
        smiles, return_tensors="pt", max_length=512,
        truncation=True,
    ).input_ids

    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_length=max_length,
            num_beams=5,          # beam search for fluency
            early_stopping=True,
        )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)


def caption_to_smiles(description: str, model, tokenizer,
                      num_candidates: int = 10) -> list[str]:
    """Generate SMILES candidates from a text description.

    Returns multiple candidates ranked by beam score.
    """
    input_ids = tokenizer(
        description, return_tensors="pt", max_length=512,
        truncation=True,
    ).input_ids

    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_length=256,
            num_beams=num_candidates,
            num_return_sequences=num_candidates,
            early_stopping=True,
        )
    candidates = [
        tokenizer.decode(ids, skip_special_tokens=True)
        for ids in output_ids
    ]
    return candidates


# Example usage
model, tokenizer = load_molt5("laituan245/molt5-base")

# Molecule captioning
aspirin_smiles = "CC(=O)Oc1ccccc1C(=O)O"
caption = smiles_to_caption(aspirin_smiles, model, tokenizer)
print(f"Caption: {caption}")
# Output: "The molecule is a member of the class of benzoic acids
#          that is salicylic acid in which the hydrogen of the
#          hydroxy group has been replaced by an acetyl group."

# Text-conditional generation
target_desc = "a nonsteroidal anti-inflammatory drug with analgesic properties"
candidates = caption_to_smiles(target_desc, model, tokenizer,
                                num_candidates=5)
for i, smi in enumerate(candidates):
    print(f"  Candidate {i+1}: {smi}")

A critical step after text-conditional generation is validity filtering. Not every SMILES string generated by MolT5 represents a valid molecule. RDKit provides the filter:

from rdkit import Chem
from rdkit.Chem import Descriptors, Draw


def filter_valid_molecules(smiles_list: list[str]) -> list[dict]:
    """Filter generated SMILES for chemical validity.

    Returns valid molecules with basic property annotations.
    """
    valid = []
    for smi in smiles_list:
        mol = Chem.MolFromSmiles(smi)
        if mol is not None:
            valid.append({
                "smiles": Chem.MolToSmiles(mol),  # canonicalize
                "molecular_weight": Descriptors.MolWt(mol),
                "logp": Descriptors.MolLogP(mol),
                "num_h_donors": Descriptors.NumHDonors(mol),
                "num_h_acceptors": Descriptors.NumHAcceptors(mol),
                "valid": True,
            })
    return valid


# Filter candidates from MolT5 generation
valid_mols = filter_valid_molecules(candidates)
print(f"Valid: {len(valid_mols)}/{len(candidates)}")
for mol_info in valid_mols:
    print(f"  {mol_info['smiles']}  MW={mol_info['molecular_weight']:.1f}")

2. MoleculeSTM: Contrastive Molecule-Text Alignment

While MolT5 treats molecule-text translation as a sequence-to-sequence problem, MoleculeSTM (Liu et al., 2023) takes the contrastive alignment approach from Section 28.1. It trains a molecular encoder (graph neural network (GNN) operating on 2D molecular graphs) and a text encoder (SciBERT, where SciBERT is a BERT language model pretrained on 1.14 million scientific papers from Semantic Scholar) with InfoNCE loss on 281,000 molecule-description pairs from PubChem. The result is a shared embedding space where molecules and their descriptions are neighbors. (As of 2024, newer contrastive molecule-text models such as MolFM and 3D-MoLM extend this approach to incorporate 3D conformer geometry alongside 2D graphs and text, improving retrieval on structure-sensitive queries.)

MoleculeSTM's advantage over MolT5 is retrieval. Instead of generating a SMILES string from scratch (which may be invalid), MoleculeSTM retrieves the closest molecules from an existing database. Given a text query, it encodes the text, then finds the nearest molecular embeddings in the shared space. Because every result comes from the database, retrieval ensures that all returned structures are known, catalogued molecules (though their relevance to the query still depends on the quality of the learned embeddings).

import torch
import torch.nn.functional as F
import numpy as np


class MoleculeSTMRetriever:
    """Text-to-molecule retrieval using MoleculeSTM embeddings.

    Assumes pre-computed molecular embeddings stored in a matrix.
    Text queries are encoded on-the-fly and compared to the
    molecular index via cosine similarity.
    """

    def __init__(self, mol_embeddings: np.ndarray,
                 mol_smiles: list[str],
                 text_encoder, text_tokenizer):
        """
        Parameters
        ----------
        mol_embeddings : np.ndarray
            Pre-computed molecular embeddings, shape (N, D).
        mol_smiles : list[str]
            SMILES strings corresponding to each row.
        text_encoder : nn.Module
            Text encoder from the aligned MoleculeSTM model.
        text_tokenizer : tokenizer
            Tokenizer for the text encoder.
        """
        # Normalize for cosine similarity
        norms = np.linalg.norm(mol_embeddings, axis=1, keepdims=True)
        self.mol_index = mol_embeddings / norms
        self.smiles = mol_smiles
        self.text_encoder = text_encoder
        self.tokenizer = text_tokenizer

    def encode_query(self, text: str) -> np.ndarray:
        """Encode a text query to the shared embedding space."""
        inputs = self.tokenizer(
            text, return_tensors="pt", max_length=256,
            truncation=True, padding=True,
        )
        with torch.no_grad():
            hidden = self.text_encoder(**inputs).last_hidden_state
            pooled = hidden[:, 0, :]  # CLS token (first token, used as the sequence-level representation)
        z = F.normalize(pooled, dim=-1)
        return z.cpu().numpy()

    def retrieve(self, query: str, top_k: int = 10) -> list[dict]:
        """Retrieve top-k molecules matching a text description.

        Returns SMILES, similarity scores, and ranks.
        """
        z_query = self.encode_query(query)  # (1, D)
        similarities = (z_query @ self.mol_index.T).squeeze()

        top_indices = np.argsort(-similarities)[:top_k]

        results = []
        for rank, idx in enumerate(top_indices):
            results.append({
                "rank": rank + 1,
                "smiles": self.smiles[idx],
                "similarity": float(similarities[idx]),
            })
        return results


# Example: retrieve kinase inhibitors by text query
# retriever = MoleculeSTMRetriever(mol_embs, smiles_db,
#                                   text_enc, text_tok)
# hits = retriever.retrieve(
#     "a selective EGFR kinase inhibitor with quinazoline scaffold",
#     top_k=5,
# )
# for hit in hits:
#     print(f"  #{hit['rank']} sim={hit['similarity']:.3f} {hit['smiles']}")

MoleculeSTM also supports text-guided molecular editing: given a starting molecule and a text instruction ("make it more water-soluble"), the model adjusts the molecular embedding in the direction indicated by the text and retrieves the nearest molecule to the adjusted vector. This is analogous to the arithmetic in word2vec (king - man + woman = queen) but operating across modalities.

def text_guided_edit(retriever: MoleculeSTMRetriever,
                     source_smiles: str,
                     edit_instruction: str,
                     alpha: float = 0.5,
                     top_k: int = 5) -> list[dict]:
    """Edit a molecule by shifting its embedding toward a text direction.

    Parameters
    ----------
    source_smiles : str
        Starting molecule SMILES.
    edit_instruction : str
        Natural language edit (e.g., "increase water solubility").
    alpha : float
        Interpolation weight: 0 = keep original, 1 = full text direction.
    """
    # Encode source molecule
    z_mol = retriever.mol_index[
        retriever.smiles.index(source_smiles)
    ]  # (D,)

    # Encode the edit instruction
    z_text = retriever.encode_query(edit_instruction).squeeze()  # (D,)

    # Interpolate: shift molecule embedding toward text
    z_edited = (1 - alpha) * z_mol + alpha * z_text
    z_edited = z_edited / np.linalg.norm(z_edited)  # re-normalize

    # Retrieve nearest molecules to the edited vector
    similarities = z_edited @ retriever.mol_index.T
    top_indices = np.argsort(-similarities)[:top_k]

    results = []
    for rank, idx in enumerate(top_indices):
        results.append({
            "rank": rank + 1,
            "smiles": retriever.smiles[idx],
            "similarity": float(similarities[idx]),
        })
    return results

3. ProteinChat: Conversational Protein Understanding

The molecule-text models above let researchers describe and retrieve small molecules in natural language; the same principle extends naturally to proteins, where free-text descriptions of function and mechanism are equally central to how biologists think.

ProteinChat (Guo et al., 2023) adapts the visual instruction tuning paradigm (LLaVA-style) to proteins. It aligns a protein structure encoder with a large language model (LLM) so that a user can ask natural language questions about a protein's 3D structure and receive informed answers.

The Three-Part Architecture

Protein encoder. ESM-2 or a 3D structure encoder (e.g., GVP-GNN, where GVP-GNN is a Geometric Vector Perceptron graph neural network that operates directly on 3D atomic coordinates and edge orientations) produces per-residue embeddings from the protein's sequence or coordinates.

Projection layer. A linear or multilayer perceptron (MLP) projection maps the protein embeddings into the LLM's input space, analogous to the visual projection in LLaVA.

Language model. A pretrained LLM (Vicuna, LLaMA) generates answers conditioned on both the projected protein embeddings and the user's question text. (As of 2024, successors such as ProteinGPT and InstructProtein have extended this paradigm with larger backbone LLMs, richer instruction-tuning datasets, and support for multi-turn dialogue about protein function and design.)

import torch
import torch.nn as nn
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer


class ProteinChatModel(nn.Module):
    """Simplified ProteinChat architecture.

    Aligns a protein encoder with an LLM via a projection layer
    for conversational protein understanding.
    """

    def __init__(
        self,
        protein_encoder_name: str = "facebook/esm2_t33_650M_UR50D",
        llm_name: str = "meta-llama/Llama-2-7b-chat-hf",
        projection_dim: int = 4096,
    ):
        super().__init__()
        self.protein_encoder = AutoModel.from_pretrained(
            protein_encoder_name
        )
        self.llm = AutoModelForCausalLM.from_pretrained(llm_name)
        self.llm_tokenizer = AutoTokenizer.from_pretrained(llm_name)

        # Freeze both backbone models
        for p in self.protein_encoder.parameters():
            p.requires_grad = False
        for p in self.llm.parameters():
            p.requires_grad = False

        # Trainable projection: protein space -> LLM input space
        esm_dim = self.protein_encoder.config.hidden_size
        self.protein_projection = nn.Sequential(
            nn.Linear(esm_dim, projection_dim),
            nn.GELU(),
            nn.Linear(projection_dim, self.llm.config.hidden_size),
        )

    def encode_protein(self, protein_ids: torch.Tensor,
                       attention_mask: torch.Tensor) -> torch.Tensor:
        """Encode protein sequence and project to LLM space.

        Returns per-residue embeddings in the LLM's input dimension.
        """
        with torch.no_grad():
            esm_output = self.protein_encoder(
                input_ids=protein_ids,
                attention_mask=attention_mask,
            )
        # Pool: mean over residue positions
        hidden = esm_output.last_hidden_state
        pooled = (hidden * attention_mask.unsqueeze(-1)).sum(dim=1)
        pooled = pooled / attention_mask.sum(dim=1, keepdim=True)
        # Project to LLM space
        return self.protein_projection(pooled)  # (batch, llm_dim)

    def answer_question(self, protein_ids, protein_mask,
                        question: str,
                        max_new_tokens: int = 256) -> str:
        """Answer a question about a protein.

        The protein embedding is prepended to the question tokens
        as a "visual prefix" in the LLM input.
        """
        # Encode protein
        protein_emb = self.encode_protein(protein_ids, protein_mask)
        # Shape: (1, llm_dim) -> (1, 1, llm_dim) as a prefix token
        prefix = protein_emb.unsqueeze(1)

        # Tokenize the question
        q_tokens = self.llm_tokenizer(
            question, return_tensors="pt"
        )
        q_embeds = self.llm.get_input_embeddings()(
            q_tokens.input_ids
        )

        # Concatenate: [protein_prefix, question_tokens]
        inputs_embeds = torch.cat([prefix, q_embeds], dim=1)

        with torch.no_grad():
            output = self.llm.generate(
                inputs_embeds=inputs_embeds,
                max_new_tokens=max_new_tokens,
            )
        return self.llm_tokenizer.decode(
            output[0], skip_special_tokens=True
        )

4. ESM3: Multimodal Protein Foundation Model

ESM3 (Hayes et al., 2024) represents the state of the art (circa 2024) in protein multimodal AI. Unlike ProteinChat, which bolts a protein encoder onto an existing LLM, ESM3 is natively multimodal: it jointly processes three tracks of protein information (sequence, structure, and function) within a single transformer architecture. Each track has its own tokenizer and embedding layer, but the transformer layers attend across all three tracks simultaneously.

ESM3 trains on 2.78 billion protein sequences (UniRef, where UniRef is the UniProt Reference Clusters database that groups protein sequences by similarity to reduce redundancy), 236 million predicted structures (ESMFold), and functional annotations (InterPro, where InterPro is a database that classifies protein sequences into families and predicts functional domains and sites), using a masked language modeling objective that randomly masks tokens from any combination of tracks. At inference, you provide any subset of inputs (sequence only, structure only, sequence plus partial function) and ESM3 generates the missing tracks, serving as a unified interface for protein understanding, design, and function prediction.

# ESM3 usage via the esm Python package (EvolutionaryScale)
# pip install esm

from esm.models.esm3 import ESM3
from esm.sdk.api import (
    ESMProtein, GenerationConfig, SamplingConfig
)


def protein_from_sequence(sequence: str) -> dict:
    """Use ESM3 to predict structure and function from sequence.

    Returns structure tokens and function annotations
    for a given amino acid sequence.
    """
    # Load the open-weight ESM3 model
    model = ESM3.from_pretrained("esm3_sm_open_v1")

    # Create a protein object with sequence only
    protein = ESMProtein(sequence=sequence)

    # Generate structure (predict folding from sequence)
    structure_config = GenerationConfig(
        track="structure",
        num_steps=8,           # iterative refinement steps
        temperature=0.7,
    )
    protein_with_structure = model.generate(
        protein, structure_config
    )

    # Generate function annotations
    function_config = GenerationConfig(
        track="function",
        num_steps=8,
        temperature=0.5,
    )
    protein_full = model.generate(
        protein_with_structure, function_config
    )

    return {
        "sequence": protein_full.sequence,
        "structure_tokens": protein_full.coordinates,
        "function_annotations": protein_full.function_annotations,
    }


def design_protein_for_function(function_description: str,
                                 length: int = 100) -> dict:
    """Design a protein sequence for a target function.

    Conditional generation: specify a desired function and let
    ESM3 generate a sequence predicted to realize it.
    """
    model = ESM3.from_pretrained("esm3_sm_open_v1")

    # Create protein with only function specified
    protein = ESMProtein(
        function_annotations=[function_description],
    )

    # Generate sequence conditioned on function
    seq_config = GenerationConfig(
        track="sequence",
        num_steps=16,
        temperature=0.8,
    )
    designed = model.generate(protein, seq_config)

    return {
        "designed_sequence": designed.sequence,
        "confidence": designed.confidence_scores,
    }

5. Cross-Referencing Molecular and Protein Embeddings

With separate models now aligning molecules to text and proteins to text, a natural next question arises: can we connect molecules and proteins to each other through the text modality they share? As illustrated in Figure 28.2, the shared text embedding space enables exactly this kind of emergent cross-modal matching.

Molecules interact with proteins: drugs bind to target proteins, metabolites serve as enzyme substrates, and signaling molecules activate receptors. If you have aligned molecule-text embeddings (from MoleculeSTM) and aligned protein-text embeddings (from ProteinChat or ESM3 plus a text projection), you can cross-reference through the shared text modality. A molecule's text description and a protein's function annotation, when both are nearby in text embedding space, suggest a potential interaction. This works because textual co-occurrence (for example, a drug and its target protein frequently appearing in the same sentences) correlates with biological interaction, even though cosine similarity in text space is only a rough proxy for binding affinity or functional relevance. Figure 28.2.1 illustrates cross-modal molecule-protein matching via shared text embedding space.

import numpy as np
from dataclasses import dataclass


@dataclass
class CrossModalMatcher:
    """Match molecules to proteins via shared text embeddings.

    Uses the text modality as a bridge: molecules aligned to text
    via MoleculeSTM, proteins aligned to text via ProteinChat.
    The overlapping text space enables cross-modal matching.
    """
    mol_text_embeddings: np.ndarray   # (N_mol, D)
    protein_text_embeddings: np.ndarray  # (N_prot, D)
    mol_smiles: list[str]
    protein_ids: list[str]

    def __post_init__(self):
        # Normalize for cosine similarity
        self.mol_normed = self.mol_text_embeddings / np.linalg.norm(
            self.mol_text_embeddings, axis=1, keepdims=True
        )
        self.prot_normed = self.protein_text_embeddings / np.linalg.norm(
            self.protein_text_embeddings, axis=1, keepdims=True
        )

    def find_targets(self, smiles: str,
                     top_k: int = 5) -> list[dict]:
        """Find proteins likely to interact with a molecule.

        Uses text-space similarity as a proxy for binding affinity.
        """
        mol_idx = self.mol_smiles.index(smiles)
        z_mol = self.mol_normed[mol_idx]  # (D,)

        # Cross-modal similarity
        scores = z_mol @ self.prot_normed.T
        top_indices = np.argsort(-scores)[:top_k]

        return [
            {
                "protein_id": self.protein_ids[idx],
                "similarity": float(scores[idx]),
                "rank": rank + 1,
            }
            for rank, idx in enumerate(top_indices)
        ]

    def find_ligands(self, protein_id: str,
                     top_k: int = 5) -> list[dict]:
        """Find molecules likely to bind a protein target."""
        prot_idx = self.protein_ids.index(protein_id)
        z_prot = self.prot_normed[prot_idx]

        scores = z_prot @ self.mol_normed.T
        top_indices = np.argsort(-scores)[:top_k]

        return [
            {
                "smiles": self.mol_smiles[idx],
                "similarity": float(scores[idx]),
                "rank": rank + 1,
            }
            for rank, idx in enumerate(top_indices)
        ]

# MolT5: 3 lines to load, 2 to run inference
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
tok = AutoTokenizer.from_pretrained("laituan245/molt5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("laituan245/molt5-base")
out = model.generate(tok("CC(=O)Oc1ccccc1C(=O)O", return_tensors="pt").input_ids)
print(tok.decode(out[0], skip_special_tokens=True))

# ESM-2 protein embeddings: 3 lines
from transformers import EsmModel, EsmTokenizer
esm_tok = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
esm_model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")
# Then: esm_model(**esm_tok(sequence, return_tensors="pt"))