scripts/e2e_v2_endpoint.py

"""End-to-end check: load v2 weights into the live server and call the real
/api/go_terms handler (thresholds, dynamic cap, taxon calibration all applied)."""
import asyncio, sys, ssl, urllib.request
from pathlib import Path
import torch
BASE = Path(__file__).parent.parent
sys.path.insert(0, str(BASE))
import server  # noqa

CASES = [("uricase","Q00511","fungus"),("catalase","P04040","mammal"),
         ("SOD1","P00441","mammal"),("hemoglobin","P69905","mammal"),
         ("p53","P04637","mammal"),("insulin","P01308","mammal")]

def fetch(acc):
    try:
        import certifi; ctx=ssl.create_default_context(cafile=certifi.where())
    except Exception: ctx=ssl._create_unverified_context()
    u=f"https://rest.uniprot.org/uniprotkb/{acc}.fasta"
    t=urllib.request.urlopen(u,timeout=30,context=ctx).read().decode()
    return "".join(l for l in t.splitlines() if not l.startswith(">"))

async def main():
    which = sys.argv[1] if len(sys.argv)>1 else "unified_35M_v2_enzyme.pth"
    async with server.lifespan(server.app):
        ck=torch.load(BASE/which,map_location="cpu",weights_only=False)
        server.model.load_state_dict(ck["model"]); server.model.eval()
        print(f"loaded {which}\n")
        from server import GoTermsRequest
        for name,acc,tax in CASES:
            seq=fetch(acc)
            req=GoTermsRequest(sequence=seq, taxon=tax, top_k=8)
            res=await server.get_go_terms(req)
            preds=res.get("predictions",[])
            shown=", ".join(f"{p['go_id']}:{server.go_map.get(p['go_id'],'')[:22]}({p['prob']})" for p in preds[:6])
            print(f"--- {name} ({acc}, {tax}) ood={res.get('ood')} n={res.get('n_predicted')}")
            print(f"    {shown}\n")

if __name__=="__main__":
    asyncio.run(main())