Export a trained DeepPurpose model to TorchScript or ONNX, stream a chunked CSV through batch prediction, or wrap it in a FastAPI predict endpoint.

DeepPurpose models can be deployed for production serving using various export formats. This guide covers model serialization and deployment patterns.

Deployment options

Format	Use Case	Pros	Cons
Native PyTorch	Development, research	Full flexibility	Requires PyTorch runtime
TorchScript	Production serving	No Python dependency	Limited dynamic behavior
ONNX	Cross-platform	Broad runtime support	May not support all encoders

Export to TorchScript

TorchScript enables deployment without Python runtime:

import torch
from DeepPurpose import DTI as models

model = models.model_pretrained(path_dir="outputs/model")

# Trace the model with example inputs
# Note: Tracing works best with fingerprint-based encoders (Morgan, Pubchem)
dummy_drug = torch.randn(1, model.model_drug.output_dim)
dummy_target = torch.randn(1, model.model_protein.output_dim)

# Use trace for models with static control flow
scripted = torch.jit.trace(model.model, (dummy_drug, dummy_target))
scripted.save("outputs/model_traced.pt")

# Load and run inference
loaded_model = torch.jit.load("outputs/model_traced.pt")
prediction = loaded_model(drug_embedding, target_embedding)

Note: Graph neural network encoders (MPNN, DGL_*) may have limited TorchScript support due to dynamic graph operations.

Export to ONNX

import torch
from deeppurpose import models

model = models.model_pretrained("outputs/model.pt", device="cuda")

dummy_drug = torch.randn(1, model.drug_encoder.output_size).cuda()
dummy_target = torch.randn(1, model.target_encoder.output_size).cuda()

torch.onnx.export(
    model.network,
    (dummy_drug, dummy_target),
    "outputs/model.onnx",
    input_names=["drug", "target"],
    output_names=["score"],
    opset_version=17,
    dynamic_axes={"drug": {0: "batch"}, "target": {0: "batch"}, "score": {0: "batch"}},
)

Batch-serving job pattern

For large-scale virtual screening, use batch processing:

import pandas as pd
from DeepPurpose import DTI as models
from DeepPurpose import utils

# Load model once
model = models.model_pretrained(path_dir="outputs/model")

def process_batch(drug_smiles, target_sequences, drug_encoding, target_encoding):
    """Process a batch of drug-target pairs."""
    data, _, _ = utils.data_process(
        X_drug=drug_smiles,
        X_target=target_sequences,
        y=[0] * len(drug_smiles),
        drug_encoding=drug_encoding,
        target_encoding=target_encoding,
        split_method='no_split'
    )
    return model.predict(data)

# Stream large input files
chunk_size = 10000
results = []

for chunk in pd.read_csv("inputs/predict.csv", chunksize=chunk_size):
    scores = process_batch(
        chunk["smiles"].tolist(),
        chunk["protein"].tolist(),
        drug_encoding="Morgan",
        target_encoding="CNN"
    )
    chunk["predicted_affinity"] = scores
    results.append(chunk)

pd.concat(results).to_csv("outputs/predictions.csv", index=False)

For command-line batch serving:

python scripts/serve_batch.py \
  --model outputs/model \
  --input inputs/predict.csv \
  --output outputs/predictions.csv \
  --batch-size 1024 \
  --drug-encoding Morgan \
  --target-encoding CNN

REST API serving

Wrap the model in a simple Flask/FastAPI service:

from fastapi import FastAPI
from pydantic import BaseModel
from DeepPurpose import DTI as models
from DeepPurpose import utils

app = FastAPI()
model = models.model_pretrained(path_dir="outputs/model")

class PredictionRequest(BaseModel):
    smiles: str
    protein_sequence: str

class PredictionResponse(BaseModel):
    predicted_affinity: float

@app.post("/predict", response_model=PredictionResponse)
def predict(request: PredictionRequest):
    data, _, _ = utils.data_process(
        X_drug=[request.smiles],
        X_target=[request.protein_sequence],
        y=[0],
        drug_encoding="Morgan",
        target_encoding="CNN",
        split_method='no_split'
    )
    score = model.predict(data)[0]
    return PredictionResponse(predicted_affinity=float(score))

Run with: uvicorn serve:app --host 0.0.0.0 --port 8000

Metadata to keep

Model format (ONNX vs. TorchScript) and opset if ONNX.
Encoder types (drug_encoding, target_encoding) used during training.
DeepPurpose version, PyTorch version, and CUDA version if GPU-based.
Training dataset and split method for reproducibility.
Git commit hash for the training code.

Serving & deployment