Spaces:
Paused
Paused
| #!/usr/bin/env python | |
| from __future__ import annotations | |
| import json | |
| import os | |
| from functools import lru_cache | |
| from pathlib import Path | |
| from typing import Any | |
| from assaylens.agent.model_predictor import predictor_from_artifacts | |
| from assaylens.agent.next_assay import recommend_next_assay | |
| DEFAULT_ARTIFACT_REPO = "lighteternal/assaylens-artifacts" | |
| DEMO_CASES = [ | |
| { | |
| "label": "7TM agonist signal needing pathway contrast", | |
| "request": { | |
| "compound": {"compound_id": "demo_cmpd_7tm_agonist", "name": "demo agonist", "smiles": "CCN(CC)CCO"}, | |
| "target": { | |
| "target_id": "demo_tgt_7tm", | |
| "gene": "ADRB2", | |
| "target_class": "7TM", | |
| "uniprot_id": "P07550", | |
| }, | |
| "observed_result": { | |
| "mode": "Agonist", | |
| "assay__mechanism": "G Protein Activation", | |
| "assay__detailed_mechanism": "cAMP", | |
| "assay__technology": "Luminescence", | |
| "outcome_is_active": True, | |
| "outcome_potency_pxc50": 6.8, | |
| "outcome_max_activity": 82.0, | |
| "is_quantified": True, | |
| "frequency_flag": False, | |
| "viability_flag": False, | |
| }, | |
| "available_assay_frames": [ | |
| {"mode": "Agonist", "assay__mechanism": "Barr2 Recruitment", "assay__technology": "Unknown"}, | |
| {"mode": "Antagonist", "assay__mechanism": "G Protein Activation", "assay__technology": "Unknown"}, | |
| {"mode": "Binding", "assay__mechanism": "Competition Binding", "assay__technology": "TR-FRET"}, | |
| ], | |
| }, | |
| }, | |
| { | |
| "label": "Kinase binding signal needing orthogonal confirmation", | |
| "request": { | |
| "compound": {"compound_id": "demo_cmpd_kinase", "name": "demo kinase ligand", "smiles": "CC1=NC=CC=C1"}, | |
| "target": { | |
| "target_id": "demo_tgt_kinase", | |
| "gene": "ABL1", | |
| "target_class": "Kinase", | |
| "uniprot_id": "P00519", | |
| }, | |
| "observed_result": { | |
| "mode": "Binding", | |
| "assay__mechanism": "Competition Binding", | |
| "assay__detailed_mechanism": "competition", | |
| "assay__technology": "TR-FRET", | |
| "outcome_is_active": True, | |
| "outcome_potency_pxc50": 7.1, | |
| "outcome_max_activity": 65.0, | |
| "is_quantified": True, | |
| "frequency_flag": True, | |
| "viability_flag": False, | |
| }, | |
| "available_assay_frames": [ | |
| {"mode": "Inhibition", "assay__mechanism": "Enzyme Activity", "assay__technology": "Unknown"}, | |
| {"mode": "Binding", "assay__mechanism": "Thermal Shift", "assay__technology": "Unknown"}, | |
| {"mode": "Binding", "assay__mechanism": "Competition Binding", "assay__technology": "TR-FRET"}, | |
| ], | |
| }, | |
| }, | |
| ] | |
| def default_request() -> dict[str, Any]: | |
| path = Path("examples/tiny_observations.json") | |
| if not path.exists(): | |
| return {} | |
| payload = json.loads(path.read_text(encoding="utf-8")) | |
| return payload[0] if isinstance(payload, list) and payload else payload | |
| def demo_cases() -> list[dict[str, Any]]: | |
| cases = [{"label": "Viability-confounded 7TM antagonist", "request": default_request()}] | |
| for case in DEMO_CASES: | |
| request = case["request"]() if callable(case["request"]) else case["request"] | |
| cases.append({"label": case["label"], "request": request}) | |
| return cases | |
| def default_request_text() -> str: | |
| return json.dumps(default_request(), indent=2) | |
| def default_candidate_frames_text() -> str: | |
| return json.dumps(default_request().get("available_assay_frames", []), indent=2) | |
| def case_by_label(label: str) -> dict[str, Any]: | |
| for case in demo_cases(): | |
| if case["label"] == label: | |
| return case["request"] | |
| return default_request() | |
| def fields_from_request(request: dict[str, Any]) -> tuple[Any, ...]: | |
| compound = request.get("compound", {}) | |
| target = request.get("target", {}) | |
| observed = request.get("observed_result", {}) | |
| return ( | |
| compound.get("compound_id", ""), | |
| compound.get("name", ""), | |
| compound.get("smiles", ""), | |
| target.get("target_id", ""), | |
| target.get("gene", ""), | |
| target.get("target_class", ""), | |
| target.get("uniprot_id", ""), | |
| observed.get("mode", ""), | |
| observed.get("assay__mechanism", ""), | |
| observed.get("assay__detailed_mechanism", ""), | |
| observed.get("assay__technology", ""), | |
| bool(observed.get("outcome_is_active")), | |
| bool(observed.get("is_quantified")), | |
| bool(observed.get("frequency_flag")), | |
| bool(observed.get("viability_flag")), | |
| json.dumps(request.get("available_assay_frames", []), indent=2), | |
| json.dumps(request, indent=2), | |
| ) | |
| def load_case(label: str) -> tuple[Any, ...]: | |
| return fields_from_request(case_by_label(label)) | |
| def cached_predictor(artifact_repo: str, model_artifact: str): | |
| return predictor_from_artifacts(artifact_repo, model_artifact) | |
| def cached_frame_coverage(artifact_repo: str) -> dict[str, int] | None: | |
| try: | |
| import pandas as pd | |
| from huggingface_hub import hf_hub_download | |
| path = hf_hub_download(artifact_repo, "data_profiles/frame_counts.csv", repo_type="dataset") | |
| frame = pd.read_csv(path) | |
| return {str(row["assay_frame_coarse"]): int(row["n_rows"]) for _, row in frame.iterrows()} | |
| except Exception: # noqa: BLE001 | |
| return None | |
| def run_recommendation( | |
| request_text: str, | |
| artifact_repo: str, | |
| model_artifact: str, | |
| include_llm: bool, | |
| ) -> tuple[str, list[dict[str, Any]], str, str, str]: | |
| request = json.loads(request_text) | |
| return recommend_request(request, artifact_repo, model_artifact, include_llm) | |
| def run_form_recommendation( | |
| compound_id: str, | |
| compound_name: str, | |
| smiles: str, | |
| target_id: str, | |
| gene: str, | |
| target_class: str, | |
| uniprot_id: str, | |
| mode: str, | |
| mechanism: str, | |
| detailed_mechanism: str, | |
| technology: str, | |
| active: bool, | |
| quantified: bool, | |
| frequency_flag: bool, | |
| viability_flag: bool, | |
| candidate_frames_text: str, | |
| artifact_repo: str, | |
| model_artifact: str, | |
| include_llm: bool, | |
| ) -> tuple[str, list[dict[str, Any]], str, str, str]: | |
| request = { | |
| "compound": { | |
| "compound_id": compound_id, | |
| "name": compound_name, | |
| "smiles": smiles, | |
| }, | |
| "target": { | |
| "target_id": target_id, | |
| "gene": gene, | |
| "target_class": target_class, | |
| "uniprot_id": uniprot_id, | |
| }, | |
| "observed_result": { | |
| "mode": mode, | |
| "assay__mechanism": mechanism, | |
| "assay__detailed_mechanism": detailed_mechanism, | |
| "assay__technology": technology, | |
| "outcome_is_active": active, | |
| "is_quantified": quantified, | |
| "frequency_flag": frequency_flag, | |
| "viability_flag": viability_flag, | |
| }, | |
| "available_assay_frames": json.loads(candidate_frames_text), | |
| } | |
| return recommend_request(request, artifact_repo, model_artifact, include_llm) | |
| def recommend_request( | |
| request: dict[str, Any], | |
| artifact_repo: str, | |
| model_artifact: str, | |
| include_llm: bool, | |
| ) -> tuple[str, list[dict[str, Any]], str, str, str]: | |
| predictor = cached_predictor(artifact_repo, model_artifact) | |
| frame_coverage = cached_frame_coverage(artifact_repo) | |
| result = recommend_next_assay( | |
| request, | |
| assaylens_mode="COUNTERFACTUAL_READY", | |
| frame_coverage=frame_coverage, | |
| predictor=predictor, | |
| include_llm_explanation=include_llm, | |
| ) | |
| recommendation = result.get("recommended_next_assay", {}) | |
| frame = recommendation.get("frame", {}) | |
| summary = "\n".join( | |
| [ | |
| f"Recommended frame: {frame.get('mode', 'unknown')} / {frame.get('assay__mechanism', 'unknown')} / {frame.get('assay__technology', 'unknown')}", | |
| f"Confidence: {recommendation.get('confidence', 'unknown')}", | |
| f"Coverage: {recommendation.get('coverage_status', 'unknown')} ({recommendation.get('coverage_count', 'n/a')})", | |
| f"Model source: {recommendation.get('model_source', 'unknown')}", | |
| f"Why: {recommendation.get('reason', 'unknown')}", | |
| ] | |
| ) | |
| ranked = [] | |
| for idx, candidate in enumerate(result.get("counterfactual_predictions", []), start=1): | |
| candidate_frame = candidate.get("candidate_frame", {}) | |
| ranked.append( | |
| { | |
| "rank": idx, | |
| "mode": candidate_frame.get("mode"), | |
| "mechanism": candidate_frame.get("assay__mechanism"), | |
| "technology": candidate_frame.get("assay__technology"), | |
| "active": round(float(candidate.get("predicted_active_probability", 0.0)), 4), | |
| "quantified": round(float(candidate.get("predicted_quantified_probability", 0.0)), 4), | |
| "viability": round(float(candidate.get("predicted_viability_confounding_probability", 0.0)), 4), | |
| "frequency": round(float(candidate.get("predicted_frequency_artifact_probability", 0.0)), 4), | |
| "uncertainty": round(float(candidate.get("uncertainty_score", 0.0)), 4), | |
| "distance": round(float(candidate.get("frame_distance", 0.0)), 4), | |
| "artifact_score": round(float(candidate.get("artifact_disambiguation_score", 0.0)), 4), | |
| "coverage": candidate.get("coverage_count"), | |
| "coverage_status": candidate.get("coverage_status"), | |
| "total_score": round(float(candidate.get("total_score", candidate.get("score", 0.0))), 4), | |
| "confidence": candidate.get("recommendation_confidence"), | |
| "exploratory": candidate.get("exploratory"), | |
| } | |
| ) | |
| warnings = "\n".join(f"- {warning}" for warning in result.get("warnings", [])) | |
| observed = result.get("observed_signal_interpretation", {}) | |
| evidence = "\n".join( | |
| [ | |
| f"Observed active probability: {float(observed.get('active_probability_under_observed_frame', 0.0)):.4f}", | |
| f"Viability confounding probability: {float(observed.get('viability_confounding_probability', 0.0)):.4f}", | |
| f"Frequency artifact probability: {float(observed.get('frequency_artifact_probability', 0.0)):.4f}", | |
| f"Observed-frame model source: {observed.get('model_source', 'unknown')}", | |
| ] | |
| ) | |
| return ( | |
| summary, | |
| ranked, | |
| json.dumps(result.get("agent_trace", {}), indent=2), | |
| result.get("llm_explanation") or warnings, | |
| evidence, | |
| ) | |
| def build_app(): | |
| import gradio as gr | |
| sample = default_request() | |
| compound = sample.get("compound", {}) | |
| target = sample.get("target", {}) | |
| observed = sample.get("observed_result", {}) | |
| css = """ | |
| .gradio-container { | |
| background: #f7f7f2; | |
| color: #1e2320; | |
| font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; | |
| } | |
| #assaylens-title { | |
| border-left: 6px solid #276357; | |
| padding-left: 16px; | |
| letter-spacing: 0; | |
| } | |
| .assaylens-panel { | |
| border: 1px solid #c7d0ca; | |
| background: #ffffff; | |
| border-radius: 8px; | |
| } | |
| button.primary { | |
| background: #276357 !important; | |
| color: #ffffff !important; | |
| border: 1px solid #1e2320 !important; | |
| } | |
| """ | |
| with gr.Blocks(css=css, title="AssayLens") as app: | |
| gr.Markdown( | |
| "# AssayLens\nRetrospective assay-frame evidence and next-measurement recommendation.", | |
| elem_id="assaylens-title", | |
| ) | |
| with gr.Row(equal_height=True): | |
| case_picker = gr.Dropdown( | |
| choices=[case["label"] for case in demo_cases()], | |
| value=demo_cases()[0]["label"], | |
| label="Curated case", | |
| ) | |
| load_case_button = gr.Button("Load Case") | |
| run_form = gr.Button("Run Recommendation", variant="primary") | |
| with gr.Accordion("Advanced artifact controls", open=False): | |
| with gr.Row(equal_height=True): | |
| artifact_repo = gr.Textbox(value=DEFAULT_ARTIFACT_REPO, label="Artifact repo") | |
| model_artifact = gr.Dropdown( | |
| choices=[ | |
| "models/b3_ligand_target_assay.pkl", | |
| "models/b4_ligand_target_assay_goa.pkl", | |
| ], | |
| value="models/b3_ligand_target_assay.pkl", | |
| label="Model artifact", | |
| ) | |
| include_llm = gr.Checkbox( | |
| value=False, | |
| label="Include LLM explanation when OPENAI_API_KEY is set", | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=5, elem_classes=["assaylens-panel"]): | |
| with gr.Tabs(): | |
| with gr.Tab("Structured input"): | |
| with gr.Row(): | |
| compound_id = gr.Textbox(value=compound.get("compound_id", ""), label="Compound ID") | |
| compound_name = gr.Textbox(value=compound.get("name", ""), label="Compound name") | |
| smiles = gr.Textbox(value=compound.get("smiles", ""), label="SMILES") | |
| with gr.Row(): | |
| target_id = gr.Textbox(value=target.get("target_id", ""), label="Target ID") | |
| gene = gr.Textbox(value=target.get("gene", ""), label="Gene") | |
| with gr.Row(): | |
| target_class = gr.Textbox(value=target.get("target_class", ""), label="Target class") | |
| uniprot_id = gr.Textbox(value=target.get("uniprot_id", ""), label="UniProt ID") | |
| with gr.Row(): | |
| mode = gr.Textbox(value=observed.get("mode", ""), label="Observed mode") | |
| mechanism = gr.Textbox(value=observed.get("assay__mechanism", ""), label="Mechanism") | |
| with gr.Row(): | |
| detailed_mechanism = gr.Textbox( | |
| value=observed.get("assay__detailed_mechanism", ""), | |
| label="Detailed mechanism", | |
| ) | |
| technology = gr.Textbox(value=observed.get("assay__technology", ""), label="Technology") | |
| with gr.Row(): | |
| active = gr.Checkbox(value=bool(observed.get("outcome_is_active")), label="Active") | |
| quantified = gr.Checkbox(value=bool(observed.get("is_quantified")), label="Quantified") | |
| frequency_flag = gr.Checkbox( | |
| value=bool(observed.get("frequency_flag")), | |
| label="Frequency artifact", | |
| ) | |
| viability_flag = gr.Checkbox( | |
| value=bool(observed.get("viability_flag")), | |
| label="Viability confound", | |
| ) | |
| candidate_frames = gr.Code( | |
| value=default_candidate_frames_text(), | |
| language="json", | |
| label="Candidate assay frames", | |
| ) | |
| with gr.Tab("JSON input"): | |
| request = gr.Code(value=default_request_text(), language="json", label="Observed assay request") | |
| run_json = gr.Button("Run JSON Recommendation", variant="primary") | |
| with gr.Column(scale=4, elem_classes=["assaylens-panel"]): | |
| summary = gr.Textbox(label="Recommendation", lines=4) | |
| evidence = gr.Textbox(label="Observed evidence", lines=5) | |
| ranked = gr.Dataframe(label="Ranked candidate frames") | |
| explanation = gr.Textbox(label="Warnings / LLM explanation", lines=8) | |
| trace = gr.Code(language="json", label="Agent trace") | |
| load_case_button.click( | |
| load_case, | |
| inputs=[case_picker], | |
| outputs=[ | |
| compound_id, | |
| compound_name, | |
| smiles, | |
| target_id, | |
| gene, | |
| target_class, | |
| uniprot_id, | |
| mode, | |
| mechanism, | |
| detailed_mechanism, | |
| technology, | |
| active, | |
| quantified, | |
| frequency_flag, | |
| viability_flag, | |
| candidate_frames, | |
| request, | |
| ], | |
| ) | |
| run_form.click( | |
| run_form_recommendation, | |
| inputs=[ | |
| compound_id, | |
| compound_name, | |
| smiles, | |
| target_id, | |
| gene, | |
| target_class, | |
| uniprot_id, | |
| mode, | |
| mechanism, | |
| detailed_mechanism, | |
| technology, | |
| active, | |
| quantified, | |
| frequency_flag, | |
| viability_flag, | |
| candidate_frames, | |
| artifact_repo, | |
| model_artifact, | |
| include_llm, | |
| ], | |
| outputs=[summary, ranked, trace, explanation, evidence], | |
| ) | |
| run_json.click( | |
| run_recommendation, | |
| inputs=[request, artifact_repo, model_artifact, include_llm], | |
| outputs=[summary, ranked, trace, explanation, evidence], | |
| ) | |
| return app | |
| def main() -> None: | |
| app = build_app() | |
| app.launch(server_name=os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1")) | |
| if __name__ == "__main__": | |
| main() | |