Update app.py

app.py

@@ -9,14 +9,14 @@ from rdkit import Chem
 from rdkit.Chem.Draw import rdMolDraw2D
 from huggingface_hub import hf_hub_download
 
-# 导入网络结构与预处理模块
+# 导入网络模型结构
 from IR_encoder import IRModel
 from SMILES_encoder import SmilesModel
+
+# 导入指定的统一网格预处理模块
 from ir_process import (
     preprocess_jdx_spectra_higer_500,
-    preprocess_jdx_spectra_lower_500,
-    preprocess_csv_spectra_higer_500,
-    preprocess_csv_spectra_lower_than_500
+    preprocess_csv_spectra_higer_500
 )
 
 # ==================== 全局配置 ====================
@@ -117,10 +117,10 @@ def load_model_instance(model_instance, weight_filename, device):
     return model_instance
 
 
-# ==================== 1. 模型初始化 (双分支权重) ====================
+# ==================== 1. 初始化双分支检索模型 ====================
 print("Initializing Models and Feature Libraries...")
 
-# 通用模型权重 (NIST / Custom 库)
+# 通用分支 (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 +128,7 @@ smiles_model_general = load_model_instance(
     DEVICE
 )
 
-# 专属 PS 模型权重 (PS Existing / PS Derivative 库)
+# 专属 PS 分支 (PS Existing / PS Derivative 库)
 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,7 +136,7 @@ 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")
 
@@ -153,24 +153,17 @@ ALL_LIBS = {
 }
 
 
-# ==================== 2. 置信度校准与自定义分子编码 ====================
+# ==================== 2. 置信度计算与分子编码 ====================
 def calculate_calibrated_confidence(cosine_sim: float, smiles_model_instance, model_type: str = "general") -> float:
-    """
-    使用从离线测试集拟合的 Platt Scaling 参数计算校准概率
-    """
     with torch.no_grad():
         t = torch.exp(smiles_model_instance.t_prime).cpu().item()
         b = smiles_model_instance.bias.cpu().item()
-    
-    # 1. 计算双塔模型原始 Logit
+
     raw_logit = cosine_sim * t + b
-    
-    # 2. 读取对应模型的 Platt A 与 B 参数
     params = PLATT_CALIBRATION_PARAMS.get(model_type, PLATT_CALIBRATION_PARAMS["general"])
     A = params["A"]
     B = params["B"]
-    
-    # 3. 逻辑回归推断: P = 1 / (1 + exp(-(A * Logit + B)))
+
     calibrated_logit = A * raw_logit + B
     confidence = 1.0 / (1.0 + np.exp(-calibrated_logit))
     return float(np.clip(confidence, 0.0, 1.0))
@@ -187,7 +180,7 @@ def process_custom_smiles(smiles_str, smiles_model_instance):
             original_map[can] = s
     if not valid_smiles:
         return None, [], []
-    
+
     tokenizer = smiles_model_instance.smiles_tokenizer
     encoded = [tokenizer(s, max_length=300, padding='max_length', truncation=True, return_tensors='pt') for s in valid_smiles]
     ids = torch.cat([item['input_ids'] for item in encoded], dim=0).to(DEVICE)
@@ -197,7 +190,7 @@ def process_custom_smiles(smiles_str, smiles_model_instance):
     return feats, valid_smiles, [original_map[s] for s in valid_smiles]
 
 
-# ==================== 3. 检索主任务 ====================
+# ==================== 3. 检索主业务流程 ====================
 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:
@@ -206,7 +199,7 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
     if ir_file is None:
         raise gr.Error("Please upload an IR spectrum file.")
 
-    # 动态切换模型分支
+    # 动态切换专属权重
     is_ps_lib = library_name in ["PS Existing Library", "PS Derivative Library"]
     current_ir_model = ir_model_ps if is_ps_lib else ir_model_general
     current_smiles_model = smiles_model_ps if is_ps_lib else smiles_model_general
@@ -218,34 +211,27 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
     if file_path.lower().endswith('.csv'):
         try:
             df = pd.read_csv(file_path, header=None)
-            w, i = df.iloc[:, 0].values, df.iloc[:, 1].values
-            if spectrum_type == "Transmittance":
-                ir_data = preprocess_csv_spectra_higer_500(w, i) if w[0] > 500 else preprocess_csv_spectra_lower_than_500(w, i)
-            else:
-                ir_data = preprocess_csv_spectra_higer_500(w, i)
+            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)}")
 
     elif file_path.lower().endswith('.jdx'):
         try:
             data = jcamp.jcamp_readfile(file_path)
-            w, i = np.array(data['x']), np.array(data['y'])
-            if spectrum_type == "Absorbance":
-                ir_data = preprocess_jdx_spectra_higer_500(w, i) if w[0] > 500 else preprocess_jdx_spectra_lower_500(w, i)
-            else:
-                ir_data = preprocess_jdx_spectra_higer_500(w, i)
+            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)}")
 
     if ir_data is None:
-        raise gr.Error("Unsupported file format or invalid selection. Please check your file and 'Type' setting.")
+        raise gr.Error("Unsupported file format or invalid data points. Please check your spectrum file.")
 
-    # 特征提取
+    # 光谱编码与检索
     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":
         lib_features, lib_smiles, display_smiles = process_custom_smiles(custom_smiles_str, current_smiles_model)
         if lib_features is None:
@@ -254,17 +240,16 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
         lib_features, lib_smiles = ALL_LIBS[library_name]
         display_smiles = lib_smiles
 
-    # 相似度矩阵与 Top10
     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]]
     res_scores = [sco.item() for sco in scores[0]]
 
-    # 置信度计算
+    # 置信度分析
     top1_score = res_scores[0]
     top1_calibrated_conf = calculate_calibrated_confidence(top1_score, current_smiles_model, model_type=model_type_flag)
 
-    # 渲染置信度卡片 (按要求更新说明文本)
+    # 规范置信度卡片展示
     conf_summary = (
         f'<div style="background-color: #ffffff; color: #1a1a1a; padding: 22px; border-radius: 12px; border: 1px solid #e0e0e0; text-align: left; box-shadow: 0 2px 6px rgba(0,0,0,0.03);">'
         f'<h3 style="color: #00bfa5; margin-top: 0; margin-bottom: 14px; text-align: center; font-size: 1.3em;">📊 置信度分析 (Confidence Analysis)</h3>'
@@ -276,7 +261,7 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
         f'</div>'
         f'<p style="font-size: 0.88em; color: #64748b; line-height: 1.6; margin: 0; text-align: justify;">'
         f'<b>Note:</b> 注意，该置信度校准为 manuscript 中 “Multi-perspective Interpretability Analysis” 节中的 '
-        f'<b>Confidence calibration analysis</b> 获取的拟合参数。鉴于受到搜索库大小与搜索对象的差异，置信度可能发生偏差，因此，建议根据给出的前置候选物进行综合人工考量。'
+        f'<b>Confidence calibration analysis</b> 获取的拟合参数，鉴于受到搜索库大小与搜索对象的差异，置信度可能发生偏差，因此，建议根据给出的前置候选物进行综合人工考量。'
         f'</p>'
         f'</div>'
     )
@@ -285,6 +270,10 @@ 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})
+
     return (
         gr.update(visible=True),
         grid_img_path,
@@ -294,12 +283,12 @@ def run_unified_retrieval(ir_file, spectrum_type, library_name, custom_smiles_st
             "Cosine Similarity": [round(s, 4) for s in res_scores],
             "Calibrated Probability": [f"{calculate_calibrated_confidence(s, current_smiles_model, model_type=model_type_flag)*100:.2f}%" for s in res_scores]
         }),
-        gr.update(visible=True, value=pd.DataFrame({'wavenumber': w, 'intensity': i})),
+        gr.update(visible=True, value=plot_df),
         gr.update(value=conf_summary, visible=True)
     )
 
 
-# ==================== 4. 界面前端 ====================
+# ==================== 4. 界面前端构建 ====================
 CSS = """
 * { font-family: 'Times New Roman', Times, serif !important; }
 body { background-color: #ffffff !important; color: #333333 !important; }