server_infer.py

import torch
import json
import os
import torch.nn.functional as F

# 定义默认路径常量
BASE_DIR = 'drug_target_activity/candidates'
PATH_CONFIG = {
    "protein": {
        "json": os.path.join(BASE_DIR, "unique_targets.json"),
        "pt": os.path.join(BASE_DIR, "protein_candidates.pt"),
        "target_type": "molecule" # 如果搜蛋白，目标是分子
    },
    "molecule": {
        "json": os.path.join(BASE_DIR, "unique_compounds.json"),
        "pt": os.path.join(BASE_DIR, "molecule_candidates.pt"),
        "target_type": "protein" # 如果搜分子，目标是蛋白
    }
}

def _find_id_by_query(metadata, query):
    """
    辅助函数：在元数据字典中查找 query 对应的 Unique ID (Key)。
    搜索优先级：Key (ID) -> Values (Sequence/Name/Gene...)
    """
    # 1. 尝试直接匹配 Key (UniProtID 或 SMILES)
    if query in metadata:
        return query, "Direct ID Match"

    # 2. 遍历 Values 进行匹配
    # 注意：这在大数据量下可能稍慢，但在 inference 场景可接受
    for unique_id, info in metadata.items():
        # 检查 info 字典里的所有值 (target__gene, compound__name, sequence 等)
        # 将所有值转为字符串进行比对
        for key, val in info.items():
            if val and str(val).strip() == str(query).strip():
                return unique_id, f"Matched field: {key}"
    
    return None, None

def retrieve_topk(input_type, input_query, topk=5, device='cpu'):
    """
    双塔检索主函数。
    
    Args:
        input_type (str): 'protein' 或 'molecule'
        input_query (str): 查询内容 (ID, Name, Gene, Sequence, SMILES...)
        topk (int): 返回前 K 个结果
        device (str): 计算设备
        
    Returns:
        list: 包含 TopK 结果的字典列表
    """
    input_type = input_type.lower()
    if input_type not in PATH_CONFIG:
        raise ValueError("input_type must be 'protein' or 'molecule'")

    # 获取当前侧（Query Side）和对侧（Target Side）的配置
    q_config = PATH_CONFIG[input_type]
    t_type = q_config['target_type']
    t_config = PATH_CONFIG[t_type]

    print(f"[*] Searching for {input_type}: '{input_query}'...")

    # ---------------------------------------------------------
    # 1. 解析 Query，找到唯一的 ID
    # ---------------------------------------------------------
    if not os.path.exists(q_config['json']):
        raise FileNotFoundError(f"Metadata file not found: {q_config['json']}")
        
    with open(q_config['json'], 'r', encoding='utf-8') as f:
        q_meta = json.load(f)

    query_id, match_reason = _find_id_by_query(q_meta, input_query)
    
    if not query_id:
        print(f"[!] Query not found in {input_type} metadata.")
        return []
    
    print(f"    -> Resolved to ID: {query_id}... ({match_reason})") #query_id[:20]
    print(f"    -> Info: {q_meta[query_id]}")

    # ---------------------------------------------------------
    # 2. 获取 Query Vector
    # ---------------------------------------------------------
    if not os.path.exists(q_config['pt']):
        raise FileNotFoundError(f"Vector file not found: {q_config['pt']}")
        
    q_pt_data = torch.load(q_config['pt'], map_location=device)
    
    try:
        # 在 .pt 文件的 ids 列表中找到 query_id 的索引
        q_index = q_pt_data['ids'].index(query_id)
        q_vector = q_pt_data['vectors'][q_index] # [Dim]
    except ValueError:
        print(f"[!] ID {query_id} found in JSON but NOT in .pt vector file. (Did you update vectors?)")
        return []

    # ---------------------------------------------------------
    # 3. 加载 Database Vectors (对侧)
    # ---------------------------------------------------------
    if not os.path.exists(t_config['pt']):
        raise FileNotFoundError(f"Target vector file not found: {t_config['pt']}")
    
    print(f"[*] Loading {t_type} database...")
    t_pt_data = torch.load(t_config['pt'], map_location=device)
    t_vectors = t_pt_data['vectors'] # [N, Dim]
    t_ids = t_pt_data['ids']         # List [N]

    # ---------------------------------------------------------
    # 4. 计算相似度 & TopK
    # ---------------------------------------------------------
    # 确保 Query 向量维度匹配 [1, Dim]
    q_vector = q_vector.unsqueeze(0) 
    
    # 计算余弦相似度 (假设存储的向量已经是 Normalized 的，直接点积即可)
    # score shape: [1, N]
    scores = torch.matmul(q_vector, t_vectors.T).squeeze(0)
    
    # 获取 TopK
    # k 不能超过数据库总量
    actual_k = min(topk, len(t_ids))
    top_scores, top_indices = torch.topk(scores, k=actual_k)

    # ---------------------------------------------------------
    # 5. 组装结果
    # ---------------------------------------------------------
    # 加载对侧的元数据用于展示详细信息
    with open(t_config['json'], 'r', encoding='utf-8') as f:
        t_meta = json.load(f)

    results = []
    print(f"\n{' Top ' + str(actual_k) + ' Matches ':=^50}")
    
    for rank, (score, idx) in enumerate(zip(top_scores, top_indices)):
        idx = idx.item()
        score = score.item()
        
        target_id = t_ids[idx]
        target_info = t_meta.get(target_id, {"error": "Metadata missing"})
        
        # 结果对象
        res_item = {
            "rank": rank + 1,
            "score": round(score, 4),
            "id": target_id,
            "info": target_info
        }
        results.append(res_item)
        
        # 打印简略信息
        # 如果是分子，打印名字；如果是蛋白，打印基因名
        display_name = target_info.get('compound__name') or target_info.get('target__gene') or "Unknown"
        print(f"Rank {rank+1} | Score: {score:.4f} | ID: {target_id[:15]}... | {display_name}")

    return results

# ==========================================
# 使用示例
# ==========================================
if __name__ == "__main__":
    # 示例 1: 输入 Protein Gene Name -> 找 Molecule
    # input_query 可以是 "ROS1", "Kinase", "P08922", 或者是具体的 Foldseek 序列
    # "P08922": {
    #     "target__foldseek_seq": "#####################################################################################################WAFWDFPDAAQFKTKTFTDDI####KKKWKWKDFP######DIDDIDPDRIDMDFTDHGQTWMKMKMWID######HIHDIDPIDHHHADDAQLFFWAW##WDA###FKIKTFTHDGPRGSADWPWKKKWKDFPVDIDTDIGR###DMGG#HHGQTKIKMKMWIADLQGIFFIDIDMDTY#########KWKWWWFWLKIWIDTPVC##DATQIA####SPTGWQEWEDDPVQQWIWIDGAQWIWIDRNPHNDDCPRIDTQDGDPFGWREWEAQPLQQKIWTATPQWIWMDHNVDPPDIDTQRAPPGTRFRYWYDPQAQQWIWTQDLQGIWIFGD#########IADQDGDNLFQEWDDQN##QWIWTQPLFQQFIKIAHV######TD#####D#AWQYWEAD##WIWTDNQQFIWTD####TDTRFHHNDVV#####GIHDIYMD###########EW##WDWQAAFFKIKIFTHG############################WWKKKWKWF#####DIDIDTRHRDRMDMGGGDHGQTKIWIWMWTDD###I#DIH#TDIDHHAHD#PDFKWKWFAFQQGIWIDTS######HGDDRPH#AWLDWDDDDQKIWTAGPFFWIWIWGR########HTD#VTTQFNEWEFQPQQQKIWGFG##QWIKIAHPVPRDIDTQDGH#AGWQYKYDDNQLQKIWTDHQFWIWIDGLNNPPTDTQGGA######GWAEKEDDPPVQKIWTWID####IWIKMWGHP########IDT###########HAWYDHRQKIWGQH##QFIWMDGR######G#####HTNYMDMDDDVVNDAH##DPDNFQQEFADFDPVQWDWDD###WTKIFGDGTD#RRDFQKWKFKWKD############HGIDIDG#######IWIWIFIHGSNDIHPIPIDRD#HH#DAFEKWFDKDKAKD###########KMKIFTGDTP###AD#FWKKKWKDK#############DIDIDG##DGMDIDG#####IKMKMWMATAHPHTGYDTDDMDMDGH#####FFWKFFFDFQWGFTARQVVRDGPDIQGHPGGWQEKDDDPVQCWIWTAGSQWIWIANNPVRDIDTLDGH##RPQQNYKEDQPAQQKIKTKH######IWMWIWFSFDP#TDTDTL######WYWQDKYADFQQCKIWTWTQDP#AIFIWIARPVVSDIDGQA##############DDPDD##QFQGEWDWAQLDVVFIWIWGDGFQFIKIDGPNNRDIDTQAGH########WRYWEDDDFKIWTWID####IKIWMATRVH#DT##IR#####RYMYGDDPSRDDDD#L#QAWW######KAWDDAAQFKTKIFGGDT########HDYYQKKKKKWKFWD##########IDIDIDRDRMDMDGGHHFQTKMKMWMTMDGDSD#########DIDMDGHHADAFDEFDDKDWDAQAQWKIKIFGDD#DGGSHDLQQKFKWKDKP#DDIPP#DTHDCCPDDP#GRMDMGGRHHGQDKMKMWMWIAG####RIDIYDIDMDHGHHAFAWWGSFDAEQFKTKTWGADG#####DDKFKWKFW####DTDTADWPFDD##TTTTMGGGHHAQTKMWMWMKTATPSGDIHIYDIDMDTHHFDWWAAKAQWDQD####FKTKIFDTDRR#FDWFWKWKWKFA############IDTDDIGD####DMDG#####HWMKMKMWIGGPRGIHDIYDIDHGGG######################################################################################PFAAADVVQKAFDAFP######TKTWIFGQQLVHDPPGRDIWIKGWGD#######VR#VVSQSSVQVPDDDQQAWHFSGWHPPDPRTITITHDQPQAFQLVNLQCQQP#####GLADPLNLLQQLLQVLVVLVRCVVQQKDQQQQARNQKGWNDNDSPDDIGIHRHGR####################PLLLFALCCLPGVIDHLLRVLSSSLVRSVCSLQSNDDRPPVADSVRSNVCLNVPHGGDDRDPDFFLSVVLSVLSSDNDSVSRDGSVVSSVSSVVVSVVVVVV#V########################################################################################################################",
    #     "target__class": "Kinase",
    #     "target__gene": "ROS1"
    # },
    retrieve_topk(
        input_type="protein", 
        input_query="ROS1", 
        topk=5
    )
    retrieve_topk(
        input_type="protein", 
        input_query="#####################################################################################################WAFWDFPDAAQFKTKTFTDDI####KKKWKWKDFP######DIDDIDPDRIDMDFTDHGQTWMKMKMWID######HIHDIDPIDHHHADDAQLFFWAW##WDA###FKIKTFTHDGPRGSADWPWKKKWKDFPVDIDTDIGR###DMGG#HHGQTKIKMKMWIADLQGIFFIDIDMDTY#########KWKWWWFWLKIWIDTPVC##DATQIA####SPTGWQEWEDDPVQQWIWIDGAQWIWIDRNPHNDDCPRIDTQDGDPFGWREWEAQPLQQKIWTATPQWIWMDHNVDPPDIDTQRAPPGTRFRYWYDPQAQQWIWTQDLQGIWIFGD#########IADQDGDNLFQEWDDQN##QWIWTQPLFQQFIKIAHV######TD#####D#AWQYWEAD##WIWTDNQQFIWTD####TDTRFHHNDVV#####GIHDIYMD###########EW##WDWQAAFFKIKIFTHG############################WWKKKWKWF#####DIDIDTRHRDRMDMGGGDHGQTKIWIWMWTDD###I#DIH#TDIDHHAHD#PDFKWKWFAFQQGIWIDTS######HGDDRPH#AWLDWDDDDQKIWTAGPFFWIWIWGR########HTD#VTTQFNEWEFQPQQQKIWGFG##QWIKIAHPVPRDIDTQDGH#AGWQYKYDDNQLQKIWTDHQFWIWIDGLNNPPTDTQGGA######GWAEKEDDPPVQKIWTWID####IWIKMWGHP########IDT###########HAWYDHRQKIWGQH##QFIWMDGR######G#####HTNYMDMDDDVVNDAH##DPDNFQQEFADFDPVQWDWDD###WTKIFGDGTD#RRDFQKWKFKWKD############HGIDIDG#######IWIWIFIHGSNDIHPIPIDRD#HH#DAFEKWFDKDKAKD###########KMKIFTGDTP###AD#FWKKKWKDK#############DIDIDG##DGMDIDG#####IKMKMWMATAHPHTGYDTDDMDMDGH#####FFWKFFFDFQWGFTARQVVRDGPDIQGHPGGWQEKDDDPVQCWIWTAGSQWIWIANNPVRDIDTLDGH##RPQQNYKEDQPAQQKIKTKH######IWMWIWFSFDP#TDTDTL######WYWQDKYADFQQCKIWTWTQDP#AIFIWIARPVVSDIDGQA##############DDPDD##QFQGEWDWAQLDVVFIWIWGDGFQFIKIDGPNNRDIDTQAGH########WRYWEDDDFKIWTWID####IKIWMATRVH#DT##IR#####RYMYGDDPSRDDDD#L#QAWW######KAWDDAAQFKTKIFGGDT########HDYYQKKKKKWKFWD##########IDIDIDRDRMDMDGGHHFQTKMKMWMTMDGDSD#########DIDMDGHHADAFDEFDDKDWDAQAQWKIKIFGDD#DGGSHDLQQKFKWKDKP#DDIPP#DTHDCCPDDP#GRMDMGGRHHGQDKMKMWMWIAG####RIDIYDIDMDHGHHAFAWWGSFDAEQFKTKTWGADG#####DDKFKWKFW####DTDTADWPFDD##TTTTMGGGHHAQTKMWMWMKTATPSGDIHIYDIDMDTHHFDWWAAKAQWDQD####FKTKIFDTDRR#FDWFWKWKWKFA############IDTDDIGD####DMDG#####HWMKMKMWIGGPRGIHDIYDIDHGGG######################################################################################PFAAADVVQKAFDAFP######TKTWIFGQQLVHDPPGRDIWIKGWGD#######VR#VVSQSSVQVPDDDQQAWHFSGWHPPDPRTITITHDQPQAFQLVNLQCQQP#####GLADPLNLLQQLLQVLVVLVRCVVQQKDQQQQARNQKGWNDNDSPDDIGIHRHGR####################PLLLFALCCLPGVIDHLLRVLSSSLVRSVCSLQSNDDRPPVADSVRSNVCLNVPHGGDDRDPDFFLSVVLSVLSSDNDSVSRDGSVVSSVSSVVVSVVVVVV#V########################################################################################################################", 
        topk=5
    )

    print("\n" + "-"*50 + "\n")

    # 示例 2: 输入 Molecule Name -> 找 Protein
    # input_query 可以是 "Lorlatinib", "DB12130" 或 SMILES 串
    # "C[C@H]1OC2=C(N)N=CC(=C2)C2=C(C#N)N(C)N=C2CN(C)C(=O)C2=C1C=C(F)C=C2": {
    #     "compound__name": "Lorlatinib",
    #     "compound__drugbank_id": "DB12130",
    #     "compound__cas": "1454846-35-5",
    #     "compound__unii": "OSP71S83EU",
    #     "compound__inchikey": "IIXWYSCJSQVBQM-LLVKDONJSA-N"
    # },
    retrieve_topk(
        input_type="molecule", 
        input_query="Lorlatinib", 
        topk=3
    )
    retrieve_topk(
        input_type="molecule", 
        input_query="C[C@H]1OC2=C(N)N=CC(=C2)C2=C(C#N)N(C)N=C2CN(C)C(=O)C2=C1C=C(F)C=C2", 
        topk=3
    )