"""Hugging Face Space entry point for the CKB prediction package.
The public Space serves two stable programmatic endpoints (``/predict`` and
``/warmup``) as well as a browser-first Gradio experience. Model artifacts are
loaded lazily and never retain a visitor's submitted data.
"""
from __future__ import annotations
import html
import shutil
import tempfile
import threading
import time
import uuid
from pathlib import Path
from typing import Any
from urllib.parse import quote
import gradio as gr
import pandas as pd
import spaces
from predictions.predict_single import CKBSingleSamplePredictor
MODEL_DIR = Path(__file__).resolve().parent / "predictions"
SHAP_CACHE_DIR = Path(tempfile.gettempdir()) / "ckb_space_shap"
SUPPLEMENTARY_IMAGE_DIR = Path(__file__).resolve().parent / "assets" / "supplementary"
SHAP_CACHE_TTL_SECONDS = 60 * 60
SHAP_LIGHT_CANVAS_MARKER = "ckb-force-plot-light-canvas"
SHAP_LIGHT_CANVAS_HEAD = f"""
"""
def _shap_force_gallery(result: dict[str, Any]) -> str:
"""Create the one-cluster plus 13-outcome SHAP explanation gallery."""
try:
force_plots = result["force_plots"]
cluster = result["cluster"]
cluster_label = str(cluster.get("cluster_label") or cluster["provisional_cluster_label"])
cards = [
_shap_force_plot_card(
"cluster",
f"Frozen cluster explanation · {cluster_label}",
force_plots["cluster"][cluster_label],
cluster=True,
)
]
for outcome_code, outcome in result["outcomes"].items():
cards.append(
_shap_force_plot_card(
"outcome",
f"{outcome_code} · {outcome['outcome']}",
force_plots[outcome_code],
)
)
except (KeyError, OSError, TypeError):
return (
''
"SHAP force plots could not be rendered for this request."
"
"
)
return """
"
def _locale_attrs(en: str, zh: str) -> str:
"""Return safe text attributes used by the language switcher."""
return (
f'data-locale-en="{html.escape(en, quote=True)}" '
f'data-locale-zh="{html.escape(zh, quote=True)}"'
)
def _outcome_presentation(code: str, outcome: dict[str, Any]) -> dict[str, Any]:
"""Get display-only reference data without influencing inference."""
fallback_name = str(outcome.get("outcome", code))
return OUTCOME_PRESENTATION.get(
code,
{
"en": fallback_name,
"zh": fallback_name,
"model_en": "Model",
"model_zh": "模型",
"low_5": 0.0,
"low_10": 0.0,
"high_5": 0.0,
"high_10": 0.0,
},
)
def _force_plot_frame(title_en: str, title_zh: str, plot_path: str) -> str:
"""Embed one trusted locally-generated SHAP force plot."""
_ensure_light_shap_canvas(plot_path)
plot_url = "/gradio_api/file=" + quote(Path(plot_path).resolve().as_posix(), safe="/")
aria_en = title_en + " SHAP force plot"
aria_zh = title_zh + " SHAP 力图"
return f"""
"""
def _supplementary_image_url(path: Path) -> str:
"""Return a temporary Gradio file URL for a bundled static figure."""
return "/gradio_api/file=" + quote(path.resolve().as_posix(), safe="/")
def _risk_reference_figure(
code: str,
presentation: dict[str, Any],
*,
is_high_risk: bool,
) -> str:
"""Show the supplied low/high cumulative-risk figure for this outcome."""
risk_key = "high" if is_high_risk else "low"
risk_en, risk_zh = ("High risk", "高风险") if is_high_risk else ("Low risk", "低风险")
figure_path = SUPPLEMENTARY_IMAGE_DIR / f"{code}_{risk_key}.png"
if not figure_path.is_file():
return _cumulative_risk_svg(code, presentation)
display_code = code.replace("_", "–")
alt_en = f"{display_code} cumulative risk curve with the {risk_en.lower()} trajectory emphasised"
alt_zh = f"{display_code} 累积风险曲线，突出显示{risk_zh}轨迹"
return f"""
Cluster profile illustration is unavailable.
'
return f"""
Low-risk reference
High-risk reference
{low_5:.1f}% {low_10:.1f}% {high_5:.1f}% {high_10:.1f}% 0 years 5 years 10 years
"""
def _outcome_image_card(code: str, outcome: dict[str, Any], plot_path: str) -> str:
"""Render a compact outcome card with context above its SHAP force plot."""
presentation = _outcome_presentation(code, outcome)
probability = max(0.0, min(1.0, float(outcome["probability"])))
threshold = max(0.0, min(1.0, float(outcome["youden_threshold"])))
is_high_risk = probability >= threshold
risk_en, risk_zh = ("High risk", "\u9ad8\u98ce\u9669") if is_high_risk else ("Low risk", "\u4f4e\u98ce\u9669")
decision_en = "at or above" if is_high_risk else "below"
decision_zh = "\u8fbe\u5230\u6216\u9ad8\u4e8e" if is_high_risk else "\u4f4e\u4e8e"
explanation_en = (
f"Predicted probability is {probability * 100:.1f}%, "
f"{decision_en} the frozen threshold of {threshold * 100:.1f}%. Reference cumulative incidence: "
f"low-risk {presentation['low_5']:.1f}% at 5 years and {presentation['low_10']:.1f}% at 10 years; "
f"high-risk {presentation['high_5']:.1f}% and {presentation['high_10']:.1f}%."
)
explanation_zh = (
f"\u9884\u6d4b\u6982\u7387\u4e3a {probability * 100:.1f}%\uff0c{decision_zh}\u51bb\u7ed3\u9608\u503c "
f"{threshold * 100:.1f}%\u3002\u53c2\u8003\u7d2f\u79ef\u53d1\u751f\u7387\uff1a\u4f4e\u98ce\u9669\u7ec4 5 \u5e74 {presentation['low_5']:.1f}%\u300110 \u5e74 "
f"{presentation['low_10']:.1f}%\uff1b\u9ad8\u98ce\u9669\u7ec4\u5206\u522b\u4e3a {presentation['high_5']:.1f}% \u548c {presentation['high_10']:.1f}%\u3002"
)
display_code = code.replace("_", "\u2013")
title_en, title_zh = str(presentation["en"]), str(presentation["zh"])
return f"""
{_risk_reference_figure(code, presentation, is_high_risk=is_high_risk)}
{risk_en}
{html.escape(explanation_en)}
{_force_plot_frame(display_code + ' ' + title_en, display_code + ' ' + title_zh, plot_path)}
"""
def _cluster_image_card(label: str, plot_path: str) -> str:
"""Render only the assigned cluster profile and its matching SHAP force plot."""
title_en, title_zh = f"Cluster {label}", f"\u805a\u7c7b {label}"
state_en, state_zh = "Assigned cluster", "\u672c\u6b21\u5206\u914d"
return f"""
{_cluster_profile_figure(label)}
{_force_plot_frame(title_en, title_zh, plot_path)}
"""
def _image_results_gallery(result: dict[str, Any]) -> str:
"""Create the assigned-cluster and 13-outcome image-result panels."""
try:
force_plots = result["force_plots"]
cluster = result["cluster"]
selected_label = str(cluster.get("cluster_label") or cluster["provisional_cluster_label"])
cluster_cards = [_cluster_image_card(selected_label, force_plots["cluster"][selected_label])]
outcome_cards = [
_outcome_image_card(code, outcome, force_plots[code])
for code, outcome in result["outcomes"].items()
]
except (KeyError, OSError, TypeError):
return 'Image results could not be rendered for this request.
'
return """
""" + "".join(cluster_cards) + """
""" + "".join(outcome_cards) + "
"
@spaces.GPU
def run_prediction(
features: dict[str, float],
sample_id: str | None = None,
*,
include_force_plots: bool = False,
) -> dict[str, Any]:
"""Run one prediction; browser-only SHAP files are temporary and auto-expire."""
predictor = get_predictor()
with _inference_lock:
if include_force_plots:
SHAP_CACHE_DIR.mkdir(parents=True, exist_ok=True)
_clean_expired_shap_reports()
output_dir = SHAP_CACHE_DIR / uuid.uuid4().hex
try:
result = predictor.predict_one(
features,
output_dir=output_dir,
make_force_plots=True,
sample_id=sample_id,
)
result["_shap_force_gallery_html"] = _image_results_gallery(result)
except Exception:
shutil.rmtree(output_dir, ignore_errors=True)
raise
else:
with tempfile.TemporaryDirectory(prefix="ckb_prediction_") as output_dir:
result = predictor.predict_one(
features,
output_dir=output_dir,
make_force_plots=False,
sample_id=sample_id,
)
# The predictor writes an audit JSON to its output directory. The API is
# intentionally stateless, so that temporary file is removed and its now
# invalid local path is not returned to callers.
result.pop("result_json", None)
result.pop("force_plots", None)
return result
def warm_up_model() -> str:
"""Load prediction-time artifacts after a visitor opens the demo."""
try:
predictor = get_predictor()
with _inference_lock:
predictor._ensure_cluster_shap()
for outcome in predictor.outcome_meta:
predictor._load_outcome_bundle(outcome)
return (
''
'模型已就绪。 预测可立即开始；提交数据不会被此 Space 持久保存。
"
)
except Exception:
return (
''
'预热尚未完成。 仍可提交预测，系统会在请求时自动重试加载。
"
)
def predict_json_api(request: dict[str, Any]) -> dict[str, Any]:
"""Single-object API used by external Python, JavaScript, or curl clients."""
if not isinstance(request, dict):
raise gr.Error("Request must be a JSON object.")
features = request.get("features", request)
if not isinstance(features, dict):
raise gr.Error("The 'features' field must be a JSON object.")
missing = [feature for feature in FEATURES if feature not in features]
if missing:
raise gr.Error(f"Missing required input variables: {missing}")
sample_id = request.get("sample_id")
if sample_id is not None:
sample_id = str(sample_id)
if len(sample_id) > 200:
raise gr.Error("sample_id must be no longer than 200 characters.")
return run_prediction(
{feature: features[feature] for feature in FEATURES},
sample_id=sample_id,
)
def _outcome_table(result: dict[str, Any]) -> pd.DataFrame:
rows: list[dict[str, Any]] = []
for code, outcome in result["outcomes"].items():
is_high_risk = outcome["risk_class"] == "high_risk"
rows.append(
{
"代码": code,
"预测结局": outcome["outcome"],
"概率": f"{outcome['probability'] * 100:.1f}%",
"冻结阈值": f"{outcome['youden_threshold'] * 100:.1f}%",
"风险判定": "At or above threshold" if is_high_risk else "Below threshold",
}
)
return pd.DataFrame(rows).set_axis(TABLE_COLUMNS_EN, axis="columns")
def _outcome_table_html(result: dict[str, Any]) -> str:
"""Render the outcome table as fully bilingual native HTML."""
rows: list[str] = []
for code, outcome in result["outcomes"].items():
presentation = _outcome_presentation(code, outcome)
probability = max(0.0, min(1.0, float(outcome["probability"])))
threshold = max(0.0, min(1.0, float(outcome["youden_threshold"])))
is_high_risk = probability >= threshold
risk_en, risk_zh = ("At or above threshold", "达到或高于阈值") if is_high_risk else ("Below threshold", "低于阈值")
rows.append(
f"""
{html.escape(code.replace('_', '–'))}
{html.escape(str(presentation['en']))} {probability * 100:.1f}%
{threshold * 100:.1f}%
{risk_en}
"""
)
return """
"
def _cluster_card(cluster: dict[str, Any]) -> str:
label = str(cluster.get("cluster_label") or cluster["provisional_cluster_label"])
assigned = cluster["classification_status"] == "assigned"
badge_class = "" if assigned else " ckb-result-badge--review"
badge_text = "已完成分配" if assigned else "建议复核"
return f"""
最近中心距离 {cluster['distance_nearest']:.4f}
与次近中心的间隔 {cluster['margin']:.4f}
"""
def _risk_visual(result: dict[str, Any]) -> str:
"""Render a responsive, accessible probability chart without extra dependencies."""
rows = sorted(
result["outcomes"].items(),
key=lambda item: float(item[1]["probability"]),
reverse=True,
)
bars: list[str] = []
for code, outcome in rows:
probability = max(0.0, min(1.0, float(outcome["probability"])))
threshold = max(0.0, min(1.0, float(outcome["youden_threshold"])))
is_high_risk = outcome["risk_class"] == "high_risk"
row_class = " ckb-risk-row--high" if is_high_risk else ""
accessible_label = (
f"{code}: {outcome['outcome']}, probability {probability * 100:.1f}%, "
f"threshold {threshold * 100:.1f}%"
)
bars.append(
f''
f'{html.escape(code)} '
f''
f' "
f'{probability * 100:.1f}% '
"
"
)
return """
"
def predict_for_demo(
*values: float | int | None,
) -> tuple[str, str, str, dict[str, Any], str]:
missing = [
FEATURE_LABELS[feature]
for feature, value in zip(FEATURES, values, strict=True)
if value is None
]
if missing:
raise gr.Error(f"请补全以下变量后再预测：{', '.join(missing)}")
features = {
feature: float(value)
for feature, value in zip(FEATURES, values, strict=True)
if value is not None
}
result = run_prediction(
features,
sample_id="web_demo",
include_force_plots=True,
)
shap_gallery = result.pop(
"_shap_force_gallery_html",
'Image results are unavailable for this request.
',
)
return (
_cluster_card(result["cluster"]),
_outcome_table_html(result),
_risk_visual(result),
result,
shap_gallery,
)
def restore_example_values() -> list[float]:
"""Restore the prefilled synthetic demonstration values without model work."""
return [EXAMPLE_VALUES[feature] for feature in FEATURES]
with gr.Blocks(
title="CKB 17-variable Risk Predictor",
fill_width=True,
) as demo:
with gr.Column(elem_id="ckb-shell"):
gr.HTML(
"""
跳转到预测输入
17 变量风险预测
一次完整输入，返回冻结 K=3 聚类分配与 13 项结局风险概率。表单与 API 共用同一套模型推理逻辑。
"""
)
gr.HTML(
"""
17 完整输入变量 K = 3 冻结最近中心聚类 13 结局风险概率
"""
)
warm_status = gr.HTML(
''
'正在后台预热模型。 你可以先填写表单；模型完成加载后即可提交。
",
elem_id="warm-status",
)
with gr.Column(elem_id="risk-inputs", elem_classes="ckb-panel"):
gr.HTML(
"""
"""
)
input_by_feature: dict[str, Any] = {}
with gr.Row(elem_classes="ckb-input-grid"):
for index, title, subtitle, group_features in FEATURE_GROUPS:
with gr.Column(min_width=250):
with gr.Group(elem_classes="ckb-form-card"):
gr.HTML(
f"""
"""
)
for feature in group_features:
component_kwargs = {
"label": DEFAULT_EN_FEATURE_LABELS[feature],
"value": EXAMPLE_VALUES[feature],
"elem_id": f"field-{feature}",
}
if feature in CATEGORICAL_CHOICES:
input_by_feature[feature] = gr.Dropdown(
choices=CATEGORICAL_CHOICES[feature],
filterable=False,
**component_kwargs,
)
else:
input_by_feature[feature] = gr.Number(
precision=0 if feature in INTEGER_FEATURES else None,
**component_kwargs,
)
# Components are displayed in user-friendly groups, while the
# predictor always receives its frozen feature order.
inputs = [input_by_feature[feature] for feature in FEATURES]
with gr.Row(elem_classes="ckb-actions"):
submit = gr.Button("运行风险预测", variant="primary", elem_id="predict-button")
reset = gr.Button("恢复演示数据", variant="secondary", elem_id="reset-button")
gr.HTML(
"""
"""
)
with gr.Column(elem_classes="ckb-results"):
gr.HTML(
"""
"""
)
cluster_output = gr.HTML(
'提交一组完整变量后，此处将显示聚类分配与距离信息。
',
elem_id="cluster-result",
)
risk_visual = gr.HTML(
'概率分布图将在预测完成后出现。
',
elem_id="risk-visual",
)
outcome_output = gr.HTML(
'Run a prediction to generate outcome details.
',
elem_id="outcome-table",
)
with gr.Accordion("Show image results (16)", open=False, elem_id="shap-result"):
shap_output = gr.HTML(
'Run a prediction to generate image results.
',
elem_id="shap-gallery",
)
with gr.Accordion("查看完整 JSON 响应", open=False, elem_id="raw-result"):
raw_output = gr.JSON(label="完整响应")
gr.HTML(
''
)
submit.click(
fn=predict_for_demo,
inputs=inputs,
outputs=[cluster_output, outcome_output, risk_visual, raw_output, shap_output],
api_name="demo_predict",
show_progress="full",
)
reset.click(fn=restore_example_values, outputs=inputs, api_name=False)
# A compact API-only endpoint. Keeping one JSON input makes it convenient
# for external Python, JavaScript, and curl clients to call the Space.
api_request = gr.JSON(label="request", visible=False)
api_response = gr.JSON(label="response", visible=False)
api_trigger = gr.Button(visible=False)
api_trigger.click(
fn=predict_json_api,
inputs=api_request,
outputs=api_response,
api_name="predict",
)
warmup_trigger = gr.Button(visible=False)
warmup_trigger.click(
fn=warm_up_model,
outputs=warm_status,
api_name="warmup",
show_progress="hidden",
)
demo.queue(default_concurrency_limit=1, max_size=16)
if __name__ == "__main__":
SHAP_CACHE_DIR.mkdir(parents=True, exist_ok=True)
demo.launch(
css=CUSTOM_CSS,
head=HEAD,
allowed_paths=[str(SHAP_CACHE_DIR), str(SUPPLEMENTARY_IMAGE_DIR)],
)