mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 16:23:20 +08:00
104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
"""LLM prompt 模块自检脚本
|
||
|
||
需先下载 MolT5 权重，默认从 models/molt5-base 读取，
|
||
可用环境变量 MOLT5_PATH 覆盖。若权重缺失则跳过（不计失败）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import sys
|
||
import traceback
|
||
from typing import Callable, List, Tuple
|
||
|
||
import torch
|
||
|
||
from lnp_ml.modeling.layers.llm_prompt import LLMPromptEncoder
|
||
|
||
MOLT5_PATH = os.environ.get("MOLT5_PATH", "models/molt5-base")
|
||
SMILES = ["CCO", "c1ccccc1"]
|
||
|
||
|
||
def _build(d_model: int = 256, **kw) -> LLMPromptEncoder:
|
||
return LLMPromptEncoder(d_model=d_model, model_name_or_path=MOLT5_PATH, **kw)
|
||
|
||
|
||
def test_output_shape() -> None:
|
||
enc = _build().eval()
|
||
out = enc(SMILES)
|
||
assert out.shape == (2, 256), f"out shape {out.shape} != (2, 256)"
|
||
|
||
|
||
def test_single_example() -> None:
|
||
enc = _build(d_model=128).eval()
|
||
out = enc(["CCO"])
|
||
assert out.shape == (1, 128), f"out shape {out.shape} != (1, 128)"
|
||
|
||
|
||
def test_encoder_frozen() -> None:
|
||
enc = _build()
|
||
n_trainable = sum(int(p.requires_grad) for p in enc.encoder.parameters())
|
||
assert n_trainable == 0, "默认应冻结 MolT5 encoder"
|
||
assert all(p.requires_grad for p in enc.proj_down.parameters()), "proj_down 应可训练"
|
||
|
||
|
||
def test_cache_used() -> None:
|
||
enc = _build().eval()
|
||
enc(SMILES)
|
||
assert all(s in enc._cache for s in SMILES), "冻结模式应缓存句向量"
|
||
# 第二次前向应复用缓存（不报错且形状一致）
|
||
assert enc(SMILES).shape == (2, 256)
|
||
|
||
|
||
def test_grad_flows_to_proj() -> None:
|
||
"""冻结 encoder 时，梯度应能流到可训练的 proj_down。"""
|
||
enc = _build()
|
||
enc(SMILES).sum().backward()
|
||
g = enc.proj_down[0].weight.grad
|
||
assert g is not None and torch.isfinite(g).all(), "proj_down 未收到有效梯度"
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
# Runner
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
|
||
def _collect_tests() -> List[Tuple[str, Callable[[], None]]]:
|
||
return [
|
||
(name, fn) for name, fn in globals().items()
|
||
if name.startswith("test_") and callable(fn)
|
||
]
|
||
|
||
|
||
def main() -> int:
|
||
try:
|
||
_build()
|
||
except Exception as e:
|
||
print(f"[SKIP] 无法加载 MolT5（path={MOLT5_PATH}）：{e}")
|
||
print(" 请先下载权重，或设置 MOLT5_PATH 后重试。")
|
||
return 0
|
||
|
||
tests = _collect_tests()
|
||
pass_count = 0
|
||
fail: List[Tuple[str, str]] = []
|
||
for name, fn in tests:
|
||
try:
|
||
fn()
|
||
except Exception:
|
||
fail.append((name, traceback.format_exc()))
|
||
print(f"[FAIL] {name}")
|
||
else:
|
||
pass_count += 1
|
||
print(f"[ OK ] {name}")
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f"Passed: {pass_count}/{len(tests)}")
|
||
if fail:
|
||
print(f"\n{'='*60}\nFailures:")
|
||
for name, tb in fail:
|
||
print(f"\n--- {name} ---\n{tb}")
|
||
return 1
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main()) |