Update app.py

app.py

@@ -25,11 +25,16 @@ TOKENIZER_PATH = r"tokenizer-smiles-roberta-1e_new"
 EXAMPLE_DIR = "example_files"
 RETRIEVAL_WEIGHTS_DIR = r"model_weights"
 PROCESSED_DATA_DIR = r"processed_data"
+os.makedirs(PROCESSED_DATA_DIR, exist_ok=True)
 
+# 密钥与私有仓库配置 (仅服务端可见)
 ADMIN_KEY = os.environ.get("ADMIN_KEY", "123456")
-HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 PRIVATE_DATASET_REPO = "Hsqcsu/Derivative_PS_data"
 
+# 衍生库内存缓存 (启动时保持为 None，只有密码正确时才按需加载)
+DERIVATIVE_PS_CACHE = None
+
 # ==================== Platt Calibration 拟合参数 ====================
 PLATT_CALIBRATION_PARAMS = {
     "general": {
@@ -84,18 +89,20 @@ def download_private_file(filename):
                 repo_id=PRIVATE_DATASET_REPO,
                 filename=filename,
                 repo_type="dataset",
-                token=HF_TOKEN,
+                token=HF_TOKEN if HF_TOKEN else True,
                 local_dir=PROCESSED_DATA_DIR
             )
             return path
         except Exception as e:
-            print(f"Download failed {filename}: {e}")
-            return local_path
+            print(f"[Auth Error] Failed to download {filename}: {e}")
+            return None
     return local_path
 
 
 def load_gz_features(filename, device):
     path = download_private_file(filename)
+    if path is None or not os.path.exists(path):
+        return None
     with gzip.open(path, 'rb') as f:
         features = torch.load(f, weights_only=False)
     return features.to(device)
@@ -103,6 +110,8 @@ def load_gz_features(filename, device):
 
 def load_smiles_list(filename):
     path = download_private_file(filename)
+    if path is None or not os.path.exists(path):
+        return []
     with open(path, 'r', encoding='utf-8') as f:
         return [line.strip() for line in f if line.strip()]
 
@@ -117,10 +126,10 @@ def load_model_instance(model_instance, weight_filename, device):
     return model_instance
 
 
-# ==================== 1. 初始化双分支检索模型 ====================
-print("Initializing Models and Feature Libraries...")
+# ==================== 1. 初始化模型与基础库 (启动时不加载衍生库) ====================
+print("Initializing Models and Base Feature Libraries...")
 
-# 通用分支 (NIST / Custom 库)
+# 通用分支模型
 ir_model_general = load_model_instance(IRModel(), "overall_best_top1_ir.pth", DEVICE)
 smiles_model_general = load_model_instance(
     SmilesModel(roberta_model_path=None, roberta_tokenizer_path=TOKENIZER_PATH, smiles_maxlen=300, feature_dim=768),
@@ -128,7 +137,7 @@ smiles_model_general = load_model_instance(
     DEVICE
 )
 
-# 专属 PS 分支 (PS Existing / PS Derivative 库)
+# 专属 PS 分支模型
 ir_model_ps = load_model_instance(IRModel(), "overall_best_nps_recall10_ir.pth", DEVICE)
 smiles_model_ps = load_model_instance(
     SmilesModel(roberta_model_path=None, roberta_tokenizer_path=TOKENIZER_PATH, smiles_maxlen=300, feature_dim=768),
@@ -136,24 +145,46 @@ smiles_model_ps = load_model_instance(
     DEVICE
 )
 
-# 载入分子特征库
+# 仅加载公开库
 library_nist_features = load_gz_features("library_NIST_smiles_features_fp16.gz", DEVICE)
 general_smiles_nist = load_smiles_list("NIST_smiles_can.txt")
 
-library_derivative_ps_features = load_gz_features("library_Derivative_PS_smiles_features_fp16.gz", DEVICE)
-ps_smiles_derivative = load_smiles_list("smiles_Derivative_PS.txt")
-
 library_existed_ps_features = load_gz_features("library_Existing_PS_smiles_features_fp16.gz", DEVICE)
 ps_smiles_existed = load_smiles_list("smiles_Existing_PS.txt")
 
-ALL_LIBS = {
+STATIC_LIBS = {
     "NIST": (library_nist_features, general_smiles_nist),
-    "PS Existing Library": (library_existed_ps_features, ps_smiles_existed),
-    "PS Derivative Library": (library_derivative_ps_features, ps_smiles_derivative)
+    "PS Existing Library": (library_existed_ps_features, ps_smiles_existed)
 }
 
 
-# ==================== 2. 置信度计算与分子编码 ====================
+# ==================== 2. 衍生库动态按需鉴权加载 ====================
+def get_authorized_derivative_ps_library(access_key):
+    """
+    只有输入正确密钥时，才动态向私有仓库请求并加载衍生库数据
+    """
+    global DERIVATIVE_PS_CACHE
+
+    if not access_key or access_key.strip() != ADMIN_KEY:
+        raise gr.Error("⛔ 访问被拒绝：请输入正确的 PS 衍生库访问密钥 (Access Key)。")
+
+    # 若已授权且已加载过，直接返回内存缓存
+    if DERIVATIVE_PS_CACHE is not None:
+        return DERIVATIVE_PS_CACHE
+
+    # 首次输入正确密码时触发安全下载与加载
+    print("--> 密钥验证成功，正在动态加载 PS 衍生库特征...")
+    feat = load_gz_features("library_Derivative_PS_smiles_features_fp16.gz", DEVICE)
+    smiles = load_smiles_list("smiles_Derivative_PS.txt")
+
+    if feat is None or len(smiles) == 0:
+        raise gr.Error("⛔ 衍生库加载失败：请检查 Hugging Face Space Settings 中是否已配置有效的 HF_TOKEN。")
+
+    DERIVATIVE_PS_CACHE = (feat, smiles)
+    return DERIVATIVE_PS_CACHE
+
+
+# ==================== 3. 置信度计算与分子编码 ====================
 def calculate_calibrated_confidence(cosine_sim: float, smiles_model_instance, model_type: str = "general") -> float:
     with torch.no_grad():
         t = torch.exp(smiles_model_instance.t_prime).cpu().item()
@@ -190,14 +221,10 @@ def process_custom_smiles(smiles_str, smiles_model_instance):
     return feats, valid_smiles, [original_map[s] for s in valid_smiles]
 
 
-# ==================== 3. 检索主业务流程 ====================
+# ==================== 4. 检索主业务流程 ====================
 def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_str, access_key, progress=gr.Progress()):
-    if library_name == "PS Derivative Library":
-        if access_key != ADMIN_KEY:
-            raise gr.Error("⛔ Access Denied: Incorrect Access Key for PS Derivative Library.")
-
     if ir_file is None:
-        raise gr.Error("Please upload an IR spectrum file.")
+        raise gr.Error("请上传待检索的红外光谱文件。")
 
     # 动态切换专属权重
     is_ps_lib = library_name in ["PS Existing Library", "PS Derivative Library"]
@@ -214,7 +241,7 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
             w_raw, i_raw = df.iloc[:, 0].values, df.iloc[:, 1].values
             ir_data = preprocess_csv_spectra_higer_500(w_raw, i_raw)
         except Exception as e:
-            raise gr.Error(f"Error processing CSV: {str(e)}")
+            raise gr.Error(f"CSV 解析处理失败: {str(e)}")
 
     elif file_path.lower().endswith('.jdx'):
         try:
@@ -222,24 +249,28 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
             w_raw, i_raw = np.array(data['x']), np.array(data['y'])
             ir_data = preprocess_jdx_spectra_higer_500(w_raw, i_raw)
         except Exception as e:
-            raise gr.Error(f"Error processing JDX: {str(e)}")
+            raise gr.Error(f"JDX 解析处理失败: {str(e)}")
 
     if ir_data is None:
-        raise gr.Error("Unsupported file format or invalid data points. Please check your spectrum file.")
+        raise gr.Error("文件格式不支持或光谱数据点无效，请检查文件。")
 
-    # 光谱编码与检索
-    ir_tensor = torch.tensor(ir_data, dtype=torch.float32).unsqueeze(0).to(DEVICE)
-    with torch.no_grad():
-        ir_feat = current_ir_model(ir_tensor).to(torch.float32)
-
-    if library_name == "Custom Library":
+    # 获取匹配目标库 (针对衍生库执行强密码鉴权)
+    if library_name == "PS Derivative Library":
+        lib_features, lib_smiles = get_authorized_derivative_ps_library(access_key)
+        display_smiles = lib_smiles
+    elif library_name == "Custom Library":
         lib_features, lib_smiles, display_smiles = process_custom_smiles(custom_smiles_str, current_smiles_model)
         if lib_features is None:
-            raise gr.Error("No valid Custom SMILES provided.")
+            raise gr.Error("未提供有效的自定义 SMILES 列表。")
     else:
-        lib_features, lib_smiles = ALL_LIBS[library_name]
+        lib_features, lib_smiles = STATIC_LIBS[library_name]
         display_smiles = lib_smiles
 
+    # 光谱编码与向量检索
+    ir_tensor = torch.tensor(ir_data, dtype=torch.float32).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        ir_feat = current_ir_model(ir_tensor).to(torch.float32)
+
     sims = ir_feat @ lib_features.to(DEVICE).to(torch.float32).t()
     scores, indices = sims.topk(min(10, lib_features.shape[0]), dim=1)
     res_smiles = [display_smiles[idx.item()] for idx in indices[0]]
@@ -270,7 +301,6 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
     grid_filename = "search_results_grid.png"
     grid_img_path = generate_mol_grid_custom(res_smiles, legends, grid_filename)
 
-    # 归一化绘图数据 (4000 -> 500)
     wavenumbers_plot = np.linspace(4000, 500, 3500)
     plot_df = pd.DataFrame({'wavenumber': wavenumbers_plot, 'intensity': ir_data})
 
@@ -288,7 +318,7 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
     )
 
 
-# ==================== 4. 界面前端构建 ====================
+# ==================== 5. 前端界面 ====================
 CSS = """
 * { font-family: 'Times New Roman', Times, serif !important; }
 body { background-color: #ffffff !important; color: #333333 !important; }