大分子AI演算法研究職位 · 面試準備手冊

展示三個核心優勢：最佳化背景 × PyTorch × 生物直覺。完整可執行，附詳細注釋。

pip install transformers torch botorch gpytorch scikit-learn matplotlib pandas

import pandas as pd
import numpy as np
import torch

# 使用 ProteinGym 中的 GB1 突變穩定性數據集（公開）
# 下載：https://github.com/OATML-Markslab/ProteinGym
# 簡化版：直接用含序列和 fitness 分數的 CSV

def load_data(csv_path):
    """
    CSV 格式：sequence | fitness_score
    fitness_score 越高代表熱穩定性越好
    """
    df = pd.read_csv(csv_path)
    sequences = df['sequence'].tolist()
    labels = torch.tensor(df['fitness_score'].values, dtype=torch.float32)
    return sequences, labels

# Demo 用假數據（實際使用時替換為真實 CSV）
def make_demo_data(n=100, seq_len=56):
    """生成示意用的隨機序列和隨機穩定性分數"""
    amino_acids = 'ACDEFGHIKLMNPQRSTVWY'
    sequences = [
        ''.join(np.random.choice(list(amino_acids), seq_len))
        for _ in range(n)
    ]
    # 假設 fitness 與序列中 'A'、'L' 比例正相關（簡化模擬）
    labels = torch.tensor([
        (s.count('A') + s.count('L')) / len(s) + np.random.normal(0, 0.05)
        for s in sequences
    ], dtype=torch.float32)
    return sequences, labels

sequences, labels = make_demo_data(n=200)
print(f"數據集大小: {len(sequences)} 條序列")

from transformers import EsmModel, EsmTokenizer

# 使用最小的 ESM-2 版本（8M 參數），適合本地跑
MODEL_NAME = "facebook/esm2_t6_8M_UR50D"
tokenizer = EsmTokenizer.from_pretrained(MODEL_NAME)
esm_model = EsmModel.from_pretrained(MODEL_NAME)
esm_model.eval()

def get_embeddings(sequences, batch_size=16):
    """
    輸入：蛋白質序列列表
    輸出：shape (N, 320) 的 embedding tensor
    320 = ESM-2 8M 模型的隱藏層維度
    
    用 mean pooling 把變長序列壓縮為固定維度向量
    """
    all_embeddings = []
    for i in range(0, len(sequences), batch_size):
        batch = sequences[i:i+batch_size]
        inputs = tokenizer(
            batch,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512
        )
        with torch.no_grad():
            outputs = esm_model(**inputs)
        
        # last_hidden_state: (batch, seq_len, hidden_dim)
        # attention_mask: (batch, seq_len) — 0 表示 padding
        hidden = outputs.last_hidden_state
        mask = inputs['attention_mask'].unsqueeze(-1).float()
        
        # Masked mean pooling（忽略 padding token）
        embeddings = (hidden * mask).sum(dim=1) / mask.sum(dim=1)
        all_embeddings.append(embeddings)
    
    return torch.cat(all_embeddings, dim=0)

print("正在提取 ESM-2 embedding（首次執行會下載模型約 30MB）...")
embeddings = get_embeddings(sequences)
print(f"Embedding shape: {embeddings.shape}")  # (200, 320)

import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# ── 數據分割 ──
X = embeddings.numpy()
y = labels.numpy()
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 標準化（重要：穩定 GP 和 NN 訓練）
scaler = StandardScaler()
X_train_s = torch.tensor(scaler.fit_transform(X_train), dtype=torch.float32)
X_test_s  = torch.tensor(scaler.transform(X_test), dtype=torch.float32)
y_train_t = torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)
y_test_t  = torch.tensor(y_test, dtype=torch.float32).unsqueeze(1)

# ── 模型架構 ──
class StabilityPredictor(nn.Module):
    """
    簡單的前饋網路，把 ESM-2 embedding 映射到穩定性分數
    Dropout 防止小數據集過擬合
    """
    def __init__(self, input_dim=320):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.LayerNorm(128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(64, 1)
        )
    def forward(self, x):
        return self.net(x)

model = StabilityPredictor(input_dim=X_train_s.shape[1])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
criterion = nn.MSELoss()
loader = DataLoader(TensorDataset(X_train_s, y_train_t), batch_size=32, shuffle=True)

# ── 訓練迴圈 ──
train_losses = []
for epoch in range(100):
    model.train()
    epoch_loss = 0
    for xb, yb in loader:
        optimizer.zero_grad()
        pred = model(xb)
        loss = criterion(pred, yb)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    train_losses.append(epoch_loss / len(loader))
    if (epoch + 1) % 20 == 0:
        print(f"Epoch {epoch+1:3d} | Loss: {train_losses[-1]:.4f}")

# ── 評估 ──
model.eval()
with torch.no_grad():
    y_pred = model(X_test_s).squeeze().numpy()
corr = np.corrcoef(y_pred, y_test)[0, 1]
print(f"\nPearson 相關係數 (test): {corr:.3f}")

from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_mll
from botorch.acquisition import ExpectedImprovement
from botorch.optim import optimize_acqf
from gpytorch.mlls import ExactMarginalLogLikelihood

def bayesian_optimization_loop(
    X_init,      # 初始訓練點 (n_init, d)
    y_init,      # 初始觀測值 (n_init, 1)
    bounds,      # 搜索範圍 (2, d)，通常是標準化後的 [-3, 3]
    n_iter=10    # BO 迭代輪數
):
    """
    貝葉斯最佳化核心流程：
    1. 用現有數據擬合高斯過程（代理模型）
    2. 用 Expected Improvement（EI）採集函數選下一個候選點
    3. 评估候選點（這裡用 NN 代理，實際應用中送去實驗）
    4. 更新數據，重複
    
    EI 的核心思想：平衡探索（不確定性高的區域）與利用（預測分數高的區域）
    """
    X_obs = X_init.clone()
    y_obs = y_init.clone()
    
    best_values = [y_obs.max().item()]
    
    for i in range(n_iter):
        # Step 1: 擬合 GP
        gp = SingleTaskGP(X_obs, y_obs)
        mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
        fit_gpytorch_mll(mll)
        
        # Step 2: 定義 EI 採集函數
        EI = ExpectedImprovement(model=gp, best_f=y_obs.max())
        
        # Step 3: 最佳化採集函數，找下一個候選點
        candidate, acq_value = optimize_acqf(
            EI,
            bounds=bounds,
            q=1,              # 每輪提議 1 個候選點
            num_restarts=5,   # 多起點避免局部最佳
            raw_samples=50
        )
        
        # Step 4: 模擬「評估」新候選點（真實應用中為實驗結果）
        with torch.no_grad():
            new_y = model(candidate).detach()
        
        # 更新觀測集
        X_obs = torch.cat([X_obs, candidate], dim=0)
        y_obs = torch.cat([y_obs, new_y], dim=0)
        best_values.append(y_obs.max().item())
        
        print(f"BO 第 {i+1:2d} 輪 | 當前最佳: {best_values[-1]:.4f} | EI: {acq_value.item():.4f}")
    
    return X_obs, y_obs, best_values

# 用 embedding 空間的主成分做為搜索空間（降維後更穩定）
from sklearn.decomposition import PCA
pca = PCA(n_components=10)
X_pca = pca.fit_transform(X_train_s.numpy())
X_bo = torch.tensor(X_pca[:20], dtype=torch.float64)  # BoTorch 需要 float64
y_bo = torch.tensor(y_train[:20], dtype=torch.float64).unsqueeze(1)

# 搜索範圍：標準化後各維度的 [-3, 3]
bounds = torch.stack([
    torch.full((10,), -3., dtype=torch.float64),
    torch.full((10,), 3.,  dtype=torch.float64)
])

print("開始貝葉斯最佳化...")
X_final, y_final, improvement_curve = bayesian_optimization_loop(X_bo, y_bo, bounds, n_iter=10)

# 執行 pipeline 後，outputs/ 目錄會產生 results.png
# 內含 2×2 四格圖，每格對應一個面試故事：
python run_pipeline.py --mode bo

• Latent space 中的連續移動不保證產生可折疊序列（對策：ProteinMPNN 做 decoder）
• GP 在高維（> 20D）效果下降，需先降維（PCA / VAE）
• 真實 oracle（濕實驗）很昂貴，需設計 batch BO 同時提交多個候選
• Demo 數據為模擬，真實上線需用 ProteinGym 或自家 assay 數據驗證

import torch
import torch.nn.functional as F

def reinforce_update(model, sequences, rewards, optimizer, baseline=None):
    """
    REINFORCE 算法核心更新
    sequences: 採樣到的序列（token ids 列表）
    rewards: 對應的 reward（如親和力代理模型分數）
    baseline: 方差縮減的基線值（通常用歷史 reward 的移動平均）
    """
    if baseline is None:
        baseline = rewards.mean()
    
    advantages = rewards - baseline  # 優勢函數：高於平均就鼓勵
    
    total_loss = 0
    for seq_tokens, advantage in zip(sequences, advantages):
        # 計算策略的 log 概率
        logits = model(seq_tokens[:-1])  # 預測下一個 token
        log_probs = F.log_softmax(logits, dim=-1)
        seq_log_prob = log_probs[range(len(seq_tokens)-1), seq_tokens[1:]].sum()
        
        # REINFORCE 目標：最大化 E[R] = 梯度上升
        loss = -advantage * seq_log_prob  # 負號：因為 optimizer 做梯度下降
        total_loss += loss
    
    optimizer.zero_grad()
    total_loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 防止梯度爆炸
    optimizer.step()
    
    return total_loss.item()