fix hang: offline mode, short timeouts, progress for large clean sets

annotate_clean.py

@@ -54,14 +54,17 @@ def load_token(path="~/.config/oncokb/token"):
         return p.read_text().strip()
     return os.environ.get("ONCOKB_TOKEN", "").strip()
 
-def vep_annotate(df, reference, gtf_path=None):
+def vep_annotate(df, reference, gtf_path=None, offline=False):
     """Try VEP if installed and cache exists, else Ensembl REST, else fallback."""
     import shutil
     vep_bin = shutil.which("vep")
     if vep_bin and gtf_path and Path(gtf_path).is_file():
         pass
     rows = []
-    for _, r in df.iterrows():
+    total = len(df)
+    for idx, (_, r) in enumerate(df.iterrows()):
+        if idx % 500 == 0 and total > 500:
+            print(f"  annotate {idx}/{total} ...", flush=True)
         chrom = str(r["chrom"])
         pos1 = int(r["position"]) + 1
         ref = str(r["ref"]); alt = str(r["alt"])
@@ -77,11 +80,21 @@ def vep_annotate(df, reference, gtf_path=None):
                 "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel",
             })
             continue
+        if offline:
+            rows.append({
+                "_orig_idx": r.name,
+                "Ген": "intergenic",
+                "HGVS_c": f"c.{pos1}{ref}>{alt}",
+                "HGVS_p": "p.(?)",
+                "HGVS": f"c.{pos1}{ref}>{alt} p.(?) ({hgvs_g})",
+                "Тип варианта и эффект": "SNV, missense_variant (predicted)" if len(ref)==1 and len(alt)==1 else "indel",
+            })
+            continue
         # Real hg38 - try Ensembl REST, expand all transcripts
         try:
             hgvs_ens = f"{chrom.replace('chr','')}:g.{pos1}{ref}>{alt}"
             url = f"https://rest.ensembl.org/vep/homo_sapiens/hgvs/{hgvs_ens}?content-type=application/json"
-            resp = requests.get(url, headers={"Content-Type": "application/json", "Accept": "application/json"}, timeout=15)
+            resp = requests.get(url, headers={"Content-Type": "application/json", "Accept": "application/json"}, timeout=(3, 5))
             if resp.status_code == 200:
                 j = resp.json()
                 tcs = j[0].get("transcript_consequences", []) if j and isinstance(j, list) and j[0] else []
@@ -118,22 +131,24 @@ def vep_annotate(df, reference, gtf_path=None):
         })
     return pd.DataFrame(rows)
 
-def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors"):
+def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors", offline=False):
     """Batch query OncoKB byHgvsVariant. Returns dict hgvs_g -> {amp, oncogenic}."""
-    if not token:
+    if not token or offline:
+        return {}
+    if len(hgvs_g_list) > 500:
+        print(f"  [oncokb] {len(hgvs_g_list)} variants - skipping OncoKB (use --offline to silence, or annotate subset)", flush=True)
         return {}
     headers = {"Authorization": f"Bearer {token}"}
     out = {}
-    for hgvs in hgvs_g_list:
+    for i, hgvs in enumerate(hgvs_g_list):
+        if i % 100 == 0 and len(hgvs_g_list) > 100:
+            print(f"  [oncokb] {i}/{len(hgvs_g_list)} ...", flush=True)
         try:
-            # OncoKB expects e.g. 7:g.140453136A>T
             params = {"hgvsg": hgvs, "tumorType": tumor_type}
-            resp = requests.get(ONCOKB_URL, headers=headers, params=params, timeout=10)
+            resp = requests.get(ONCOKB_URL, headers=headers, params=params, timeout=(3, 5))
             if resp.status_code == 200:
                 j = resp.json()
-                # j contains levelOfEvidence, oncogenic, etc.
                 amp = j.get("highestSensitiveLevel") or j.get("levelAssociated") or ""
-                # Map OncoKB levels to AMP Tier
                 level_map = {"LEVEL_1": "Tier I", "LEVEL_2": "Tier II", "LEVEL_3A": "Tier II", "LEVEL_3B": "Tier III", "LEVEL_4": "Tier IV", "LEVEL_R1": "Tier I", "LEVEL_R2": "Tier II"}
                 amp_tier = level_map.get(amp, amp)
                 oncogenic = j.get("oncogenic", "")
@@ -142,7 +157,7 @@ def fetch_oncokb(hgvs_g_list, token, tumor_type="All Solid Tumors"):
                 out[hgvs] = {"AMP": "", "ONCO": ""}
         except Exception:
             out[hgvs] = {"AMP": "", "ONCO": ""}
-        time.sleep(0.2)
+        time.sleep(0.05)
     return out
 
 def fetch_civic():
@@ -251,6 +266,7 @@ def main():
     ap.add_argument("--out", default=None, help="output xlsx (default <clean>.annotated.xlsx)")
     ap.add_argument("--tumor-type", default="All Solid Tumors", help="OncoKB tumor type (pan-cancer default)")
     ap.add_argument("--token", default="~/.config/oncokb/token", help="OncoKB token file or env")
+    ap.add_argument("--offline", action="store_true", help="skip Ensembl/OncoKB network calls (fast, offline)")
     args = ap.parse_args()
 
     clean_path = Path(args.clean)
@@ -258,13 +274,15 @@ def main():
     if df_clean.empty:
         print("clean_variants.csv is empty (or only header) - nothing to annotate")
         out = args.out or str(clean_path).replace(".csv", ".annotated.xlsx")
-        # Still create empty Excel with headers
         df_empty = pd.DataFrame(columns=["Ген","HGVS_c","HGVS_p","HGVS","Тип варианта и эффект","_orig_idx"])
         build_excel(df_clean, df_empty, out)
         return
 
+    if len(df_clean) > 500 and not args.offline:
+        print(f"[warn] {len(df_clean)} variants - network annotation will be slow. Use --offline for fast placeholder.", flush=True)
+
     # HGVS / gene / effect
-    df_annot = vep_annotate(df_clean, args.reference)
+    df_annot = vep_annotate(df_clean, args.reference, offline=args.offline)
 
     # OncoKB
     token = load_token(args.token)
@@ -273,7 +291,7 @@ def main():
     hgvs_g_for_oncokb = []
     for _, r in df_clean.iterrows():
         hgvs_g_for_oncokb.append(f"{str(r['chrom']).replace('chr','')}:g.{int(r['position'])+1}{r['ref']}>{r['alt']}")
-    oncokb_map = fetch_oncokb(hgvs_g_for_oncokb, token, args.tumor_type) if token else {}
+    oncokb_map = fetch_oncokb(hgvs_g_for_oncokb, token, args.tumor_type, offline=args.offline) if token else {}
 
     # Attach ACMG/AMP/ONCO to annot rows
     acmg_list = annotate_acmg(df_clean)