Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import pandas as pd | |
| import torch | |
| import numpy as np | |
| from transformers import BertModel, BertTokenizer | |
| import requests | |
| import os | |
| import pickle | |
| import hashlib | |
| import time | |
| import threading | |
| import re | |
| from datetime import datetime | |
| # ===== 页面设置 ===== | |
| st.set_page_config(page_title="多肽虚拟筛选工具", layout="wide") | |
| st.title("多肽虚拟筛选工具") | |
| # ===== 防休眠机制 ===== | |
| def keep_alive(): | |
| while True: | |
| time.sleep(300) | |
| try: | |
| _ = st.session_state.get("_", None) | |
| except Exception: | |
| pass | |
| if "alive_thread" not in st.session_state: | |
| alive_thread = threading.Thread(target=keep_alive, daemon=True) | |
| alive_thread.start() | |
| st.session_state.alive_thread = True | |
| # ===== 默认参数 ===== | |
| DEFAULTS = { | |
| "drug": "HAEGTFTSDV", | |
| "min_len": 10, | |
| "max_len": 12, | |
| "cut_rule": "胰蛋白酶 (KR)", | |
| "top_n": 10, | |
| "content_threshold": 10, | |
| } | |
| for key, value in DEFAULTS.items(): | |
| if key not in st.session_state: | |
| st.session_state[key] = value | |
| # 持久化状态 | |
| if "last_results" not in st.session_state: | |
| st.session_state.last_results = None | |
| if "last_params_hash" not in st.session_state: | |
| st.session_state.last_params_hash = "" | |
| if "running" not in st.session_state: | |
| st.session_state.running = False | |
| if "display_id" not in st.session_state: | |
| st.session_state.display_id = 0 | |
| # 缓存显示时使用的参数快照（只在成功完成筛选后才更新） | |
| if "cached_display_params" not in st.session_state: | |
| st.session_state.cached_display_params = "" | |
| if "cached_top_n" not in st.session_state: | |
| st.session_state.cached_top_n = 10 | |
| # ===== 侧边栏参数 ===== | |
| st.sidebar.header("筛选参数") | |
| def reset_params(): | |
| for key, value in DEFAULTS.items(): | |
| st.session_state[key] = value | |
| # 不重置缓存的显示参数，保留上次筛选报告的快照 | |
| st.sidebar.button("🔄 重置为默认参数", on_click=reset_params, use_container_width=True) | |
| drug = st.sidebar.text_input( | |
| "靶点序列（参照肽）", | |
| key="drug", | |
| help="请输入20种天然氨基酸字母序列（如：HAEGTFTSDV）" | |
| ) | |
| if drug: | |
| if re.fullmatch(r'[ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwy]+', drug): | |
| st.sidebar.success("✅ 序列格式正确") | |
| else: | |
| st.sidebar.error("❌ 序列包含非法字符，仅限20种天然氨基酸字母") | |
| min_len = st.sidebar.number_input("最小长度", min_value=5, max_value=50, key="min_len") | |
| max_len = st.sidebar.number_input("最大长度", min_value=5, max_value=50, key="max_len") | |
| if min_len > max_len: | |
| st.sidebar.error("❌ 最小长度不能大于最大长度") | |
| cut_rule = st.sidebar.selectbox( | |
| "酶切规则", | |
| options=[ | |
| "胰蛋白酶 (KR)", | |
| "胃蛋白酶 (FL)", | |
| "胰凝乳蛋白酶 (FYWM)", | |
| "木瓜蛋白酶 (KR)", | |
| "菠萝蛋白酶 (KRFY)", | |
| "无花果蛋白酶 (KR)", | |
| "中性蛋白酶 (FLIV)", | |
| "碱性蛋白酶 (KRFLIV)", | |
| ], | |
| key="cut_rule", | |
| ) | |
| cut_map = { | |
| "胰蛋白酶 (KR)": "KR", | |
| "胃蛋白酶 (FL)": "FL", | |
| "胰凝乳蛋白酶 (FYWM)": "FYWM", | |
| "木瓜蛋白酶 (KR)": "KR", | |
| "菠萝蛋白酶 (KRFY)": "KRFY", | |
| "无花果蛋白酶 (KR)": "KR", | |
| "中性蛋白酶 (FLIV)": "FLIV", | |
| "碱性蛋白酶 (KRFLIV)": "KRFLIV", | |
| } | |
| CUT = cut_map[cut_rule] | |
| top_n = st.sidebar.number_input("输出Top几", min_value=5, max_value=100, key="top_n") | |
| content_threshold = st.sidebar.number_input("蛋白含量阈值(%)", min_value=0, max_value=100, key="content_threshold") | |
| # ===== 辅助函数 ===== | |
| def get_cache_key(): | |
| raw = f"{drug}_{CUT}_{min_len}_{max_len}" | |
| return hashlib.md5(raw.encode()).hexdigest()[:8] | |
| def get_params_hash(): | |
| raw = f"{drug}_{min_len}_{max_len}_{cut_rule}_{top_n}_{content_threshold}" | |
| return hashlib.md5(raw.encode()).hexdigest() | |
| # ===== 模型加载 ===== | |
| def load_model(): | |
| with st.spinner("正在加载ProtBERT模型..."): | |
| model_name = "Rostlab/prot_bert" | |
| try: | |
| tok = BertTokenizer.from_pretrained(model_name, local_files_only=True) | |
| mod = BertModel.from_pretrained(model_name, local_files_only=True) | |
| except Exception: | |
| tok = BertTokenizer.from_pretrained(model_name) | |
| mod = BertModel.from_pretrained(model_name) | |
| return tok, mod | |
| # ===== 核心算法 ===== | |
| def cut(seq, cut_aa, min_l, max_l): | |
| out = [] | |
| now = "" | |
| for i, aa in enumerate(seq): | |
| now += aa | |
| if aa in cut_aa: | |
| if i + 1 >= len(seq) or seq[i + 1] != "P": | |
| if min_l <= len(now) <= max_l: | |
| out.append(now) | |
| now = "" | |
| if now and min_l <= len(now) <= max_l: | |
| out.append(now) | |
| return out | |
| def vec(p, tok, mod): | |
| s = " ".join(list(p)) | |
| x = tok(s, return_tensors="pt") | |
| with torch.no_grad(): | |
| y = mod(**x) | |
| return y.last_hidden_state.mean(dim=1).numpy().flatten() | |
| def sim(a, b): | |
| return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)) | |
| def fetch_proteins(): | |
| SPECIES_IDS = [ | |
| 3847, 9913, 9031, 8090, 7955, 8030, 3888, 4565, 3818, 9823, 3994, 4497 | |
| ] | |
| food_dict = {} | |
| for sid in SPECIES_IDS: | |
| url = f"https://rest.uniprot.org/uniprotkb/search?query=organism_id:{sid}&size=100&format=fasta" | |
| try: | |
| response = requests.get(url, timeout=60) | |
| if response.status_code == 200: | |
| lines = response.text.strip().split("\n") | |
| current_name, current_seq = "", "" | |
| for line in lines: | |
| if line.startswith(">"): | |
| if current_seq and len(current_seq) >= 50: | |
| food_dict[current_name] = current_seq | |
| parts = line[1:].split("|") | |
| entry_name = parts[1].strip() if len(parts) >= 2 else line[1:].strip()[:30] | |
| species = "" | |
| desc = parts[2].strip() if len(parts) >= 3 else "" | |
| if "OS=" in desc: | |
| species = desc.split("OS=")[-1].split()[0] | |
| current_name = f"{entry_name} ({species})" | |
| else: | |
| current_name = entry_name | |
| current_seq = "" | |
| else: | |
| current_seq += line.strip() | |
| if current_seq and len(current_seq) >= 50: | |
| food_dict[current_name] = current_seq | |
| except Exception: | |
| pass | |
| return food_dict | |
| PROTEIN_CONTENT = { | |
| "glycine max": 38, "glycine": 38, | |
| "bos taurus": 25, "bos": 25, | |
| "gallus gallus": 22, "gallus": 22, | |
| "oryzias latipes": 18, "oryzias": 18, | |
| "danio rerio": 18, "danio": 18, | |
| "salmo salar": 22, "salmo": 22, | |
| "parambassis": 16, | |
| "pisum sativum": 23, "pisum": 23, | |
| "triticum aestivum": 75, "triticum": 75, | |
| "arachis hypogaea": 28, "arachis": 28, | |
| "sus scrofa": 24, "sus": 24, | |
| "oryza sativa": 7, "oryza": 7, | |
| "avena sativa": 13, "avena": 13, | |
| } | |
| def get_content(src_name): | |
| src_lower = src_name.lower() | |
| if "(" in src_name and ")" in src_name: | |
| species_short = src_name.split("(")[-1].split(")")[0].strip().lower() | |
| for key, content in PROTEIN_CONTENT.items(): | |
| if key.lower() == species_short: | |
| return content | |
| for key, content in PROTEIN_CONTENT.items(): | |
| if key.lower() in src_lower: | |
| return content | |
| return 0 | |
| def make_filter_params(): | |
| return f"靶点：{drug.upper()} | 酶：{cut_rule.split('(')[0].strip()} | 长度：{min_len}-{max_len} | 含量阈值：≥{content_threshold}% | 时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" | |
| def display_results(res): | |
| st.session_state.display_id += 1 | |
| did = st.session_state.display_id | |
| # 使用缓存的参数快照，绝不实时读取侧边栏 | |
| st.caption(st.session_state.cached_display_params) | |
| st.subheader("结构相似度排名") | |
| df_top = pd.DataFrame(res[:st.session_state.cached_top_n]) | |
| df_top.index = range(1, len(df_top) + 1) | |
| show_cols = ["src", "p", "score"] | |
| if "content_pct" in df_top.columns: | |
| show_cols.append("content_pct") | |
| st.dataframe(df_top[show_cols], width='stretch') | |
| csv_top = df_top.to_csv(index=False).encode("utf-8-sig") | |
| st.download_button("下载Top结果CSV", csv_top, "out_top.csv", "text/csv", key=f"download_top_{did}") | |
| csv_all = pd.DataFrame(res).to_csv(index=False).encode("utf-8-sig") | |
| st.download_button("下载全部结果CSV", csv_all, "out_all.csv", "text/csv", key=f"download_all_{did}") | |
| st.subheader("物种贡献统计（全部候选肽）") | |
| species_counts = {} | |
| for x in res: | |
| s = x["src"].split("(")[-1].split(")")[0].strip() if "(" in x["src"] else "unknown" | |
| species_counts[s] = species_counts.get(s, 0) + 1 | |
| species_df = pd.DataFrame(list(species_counts.items()), columns=["物种", "候选肽数量"]) | |
| species_df = species_df.sort_values("候选肽数量", ascending=False).reset_index(drop=True) | |
| species_df.index = range(1, len(species_df) + 1) | |
| species_df["占比(%)"] = (species_df["候选肽数量"] / species_df["候选肽数量"].sum() * 100).round(1) | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.dataframe(species_df[["物种", "候选肽数量", "占比(%)"]].head(20), width='stretch') | |
| with col2: | |
| st.bar_chart(species_df.set_index("物种")["候选肽数量"]) | |
| # ===== 显示缓存结果（使用已保存的参数快照） ===== | |
| if st.session_state.last_results is not None: | |
| display_results(st.session_state.last_results) | |
| st.subheader("开始新的筛选") | |
| if st.button("🚀 开始筛选", type="primary"): | |
| if not re.fullmatch(r'[ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwy]+', drug): | |
| st.error("靶点序列包含非法字符，请输入有效的氨基酸序列（仅限20种天然氨基酸字母）") | |
| st.stop() | |
| if min_len > max_len: | |
| st.error("最小长度不能大于最大长度") | |
| st.stop() | |
| current_hash = get_params_hash() | |
| if current_hash == st.session_state.last_params_hash and st.session_state.last_results is not None: | |
| st.info("参数未修改，显示上次筛选结果") | |
| st.stop() | |
| st.session_state.running = True | |
| try: | |
| tok, mod = load_model() | |
| food_cache_file = "food_cache.pkl" | |
| if os.path.exists(food_cache_file): | |
| with open(food_cache_file, "rb") as f: | |
| FOOD = pickle.load(f) | |
| st.success(f"已从缓存加载 {len(FOOD)} 条蛋白序列") | |
| else: | |
| with st.spinner("正在从UniProt下载食用蛋白序列..."): | |
| FOOD = fetch_proteins() | |
| with open(food_cache_file, "wb") as f: | |
| pickle.dump(FOOD, f) | |
| st.success(f"已获取并缓存 {len(FOOD)} 条蛋白序列") | |
| allp = [] | |
| for name, seq in FOOD.items(): | |
| for p in cut(seq, CUT, min_len, max_len): | |
| allp.append({"src": name, "p": p}) | |
| seen = set() | |
| unip = [] | |
| for x in allp: | |
| species_short = x["src"].split("(")[-1].split(")")[0].strip() if "(" in x["src"] else x["src"] | |
| key = (species_short, x["p"]) | |
| if key not in seen: | |
| seen.add(key) | |
| unip.append(x) | |
| st.info(f"酶切并去重后共 {len(unip)} 条候选肽") | |
| cache_key = get_cache_key() | |
| score_cache_file = f"score_cache_{cache_key}.pkl" | |
| if os.path.exists(score_cache_file): | |
| with open(score_cache_file, "rb") as f: | |
| res = pickle.load(f) | |
| st.success(f"已从缓存加载 {len(res)} 条打分结果") | |
| else: | |
| progress_bar = st.progress(0) | |
| status_text = st.empty() | |
| v_drug = vec(drug.upper(), tok, mod) | |
| res = [] | |
| total = len(unip) | |
| for idx, x in enumerate(unip): | |
| v_p = vec(x["p"], tok, mod) | |
| s = sim(v_p, v_drug) | |
| res.append({"src": x["src"], "p": x["p"], "score": round(s, 4)}) | |
| if idx % 50 == 0: | |
| progress_bar.progress(min(idx / total, 1.0)) | |
| status_text.text(f"打分进度: {idx}/{total}") | |
| progress_bar.progress(1.0) | |
| status_text.text(f"打分完成: {total}/{total}") | |
| with open(score_cache_file, "wb") as f: | |
| pickle.dump(res, f) | |
| st.success(f"已完成打分并缓存") | |
| for x in res: | |
| x["content_pct"] = get_content(x["src"]) | |
| filtered_res = [x for x in res if x["content_pct"] >= content_threshold] | |
| if filtered_res: | |
| res = filtered_res | |
| res.sort(key=lambda x: x["score"], reverse=True) | |
| # ---- 关键修复：只在成功生成结果后才更新快照 ---- | |
| st.session_state.last_results = res | |
| st.session_state.last_params_hash = current_hash | |
| st.session_state.cached_display_params = make_filter_params() | |
| st.session_state.cached_top_n = top_n | |
| display_results(res) | |
| st.success("筛选完成！") | |
| except Exception as e: | |
| st.error(f"运行出错：{e}") | |
| finally: | |
| st.session_state.running = False |