Add data processing pipeline

data_processing.py

@@ -0,0 +1,164 @@
+"""
+MSigDB Data Processing Pipeline for Contrastive Pretraining.
+
+Strategy:
+1. Download full GMT files from Broad data server (no auth needed)
+2. Fetch brief descriptions from MSigDB HTML card pages (no auth needed)
+3. Build text-gene paired dataset
+
+Usage:
+    python data_processing.py
+"""
+
+import json
+import os
+import time
+import html
+import re
+import random
+import concurrent.futures
+from collections import defaultdict
+import requests
+
+BROAD_BASE = "https://data.broadinstitute.org/gsea-msigdb/msigdb/release"
+VERSION = "2024.1"
+
+HUMAN_GMTS = {"H": "h.all", "C1": "c1.all", "C2": "c2.all", "C3": "c3.all",
+              "C4": "c4.all", "C5": "c5.all", "C6": "c6.all", "C7": "c7.all", "C8": "c8.all"}
+MOUSE_GMTS = {"MH": "mh.all", "M1": "m1.all", "M2": "m2.all", "M3": "m3.all",
+              "M5": "m5.all", "M8": "m8.all"}
+
+
+def download_gmt(collection_code, species="Hs", version=VERSION):
+    gmts = HUMAN_GMTS if species == "Hs" else MOUSE_GMTS
+    prefix = gmts.get(collection_code)
+    if not prefix:
+        return []
+    filename = f"{prefix}.v{version}.{species}.symbols.gmt"
+    url = f"{BROAD_BASE}/{version}.{species}/{filename}"
+    try:
+        resp = requests.get(url, timeout=60)
+        resp.raise_for_status()
+    except Exception as e:
+        print(f"  Warning: {url}: {e}")
+        return []
+    gene_sets = []
+    for line in resp.text.strip().split("\\n"):
+        parts = line.split("\\t")
+        if len(parts) < 3:
+            continue
+        gene_sets.append({
+            "name": parts[0].strip(), "url": parts[1].strip(),
+            "genes": [g.strip() for g in parts[2:] if g.strip()],
+            "collection": collection_code,
+            "species": "human" if species == "Hs" else "mouse",
+        })
+    return gene_sets
+
+
+def download_all_gmts(output_dir="data/raw"):
+    os.makedirs(output_dir, exist_ok=True)
+    all_gs = []
+    print("Downloading human gene sets...")
+    for code in HUMAN_GMTS:
+        gs = download_gmt(code, "Hs")
+        all_gs.extend(gs)
+        print(f"  {code}: {len(gs)}")
+    print("Downloading mouse gene sets...")
+    for code in MOUSE_GMTS:
+        gs = download_gmt(code, "Mm")
+        all_gs.extend(gs)
+        print(f"  {code}: {len(gs)}")
+    with open(os.path.join(output_dir, "all_gmt_genesets.json"), "w") as f:
+        json.dump(all_gs, f)
+    print(f"Total: {len(all_gs)}")
+    return all_gs
+
+
+def fetch_description_html(name, species="human"):
+    url = f"https://www.gsea-msigdb.org/gsea/msigdb/{species}/geneset/{name}"
+    try:
+        resp = requests.get(url, timeout=15)
+        resp.raise_for_status()
+        match = re.findall(r'Brief\\s+description.*?<td[^>]*>(.*?)</td>', resp.text, re.DOTALL | re.IGNORECASE)
+        if match:
+            desc = re.sub(r'<[^>]+>', '', match[0]).strip()
+            desc = html.unescape(desc)
+            if desc and desc.lower() not in ["na", "n/a"]:
+                return desc
+    except Exception:
+        pass
+    return ""
+
+
+def fetch_descriptions_batch(gene_sets, max_workers=10, cache_path="data/raw/descriptions_cache.json"):
+    cache = {}
+    if os.path.exists(cache_path):
+        with open(cache_path) as f:
+            cache = json.load(f)
+    to_fetch = [(gs["name"], gs["species"]) for gs in gene_sets if gs["name"] not in cache]
+    print(f"Need to fetch {len(to_fetch)} descriptions ({len(cache)} cached)")
+    if to_fetch:
+        fetched = 0
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as ex:
+            futures = {ex.submit(fetch_description_html, n, s): (n, s) for n, s in to_fetch}
+            for f in concurrent.futures.as_completed(futures):
+                name, _ = futures[f]
+                cache[name] = f.result()
+                fetched += 1
+                if fetched % 500 == 0:
+                    with open(cache_path, "w") as fp:
+                        json.dump(cache, fp)
+                    print(f"  {fetched}/{len(to_fetch)}")
+        with open(cache_path, "w") as f:
+            json.dump(cache, f)
+    return cache
+
+
+def build_pairs(gene_sets, descriptions, min_genes=5, max_genes=2000):
+    pairs = []
+    for gs in gene_sets:
+        genes = gs["genes"]
+        if len(genes) < min_genes or len(genes) > max_genes:
+            continue
+        desc = descriptions.get(gs["name"], "")
+        parts = [f"[Collection: {gs['collection']}] [Species: {gs['species']}]",
+                 gs["name"].replace("_", " ")]
+        if desc:
+            parts.append(html.unescape(re.sub(r'<[^>]+>', ' ', desc)).strip())
+        text = "\\n".join(parts)
+        if len(text) < 30:
+            continue
+        pairs.append({"id": gs["name"], "text": text, "genes": genes,
+                      "n_genes": len(genes), "collection": gs["collection"],
+                      "species": gs["species"], "has_description": bool(desc)})
+    return pairs
+
+
+def split_and_save(pairs, output_dir="data/processed"):
+    os.makedirs(output_dir, exist_ok=True)
+    train_cols = {"C2", "C5", "C8", "C1", "M2", "M5", "M8", "M1"}
+    val_cols = {"C3", "C4", "M3"}
+    test_cols = {"H", "C6", "C7", "MH"}
+    splits = {"train": [], "val": [], "test": []}
+    for p in pairs:
+        c = p["collection"]
+        if c in train_cols: splits["train"].append(p)
+        elif c in val_cols: splits["val"].append(p)
+        elif c in test_cols: splits["test"].append(p)
+        else: splits["train"].append(p)
+    for name, data in splits.items():
+        path = os.path.join(output_dir, f"{name}.jsonl")
+        with open(path, "w") as f:
+            for r in data:
+                f.write(json.dumps(r) + "\\n")
+        print(f"{name}: {len(data)} pairs -> {path}")
+    return splits
+
+
+if __name__ == "__main__":
+    all_gs = download_all_gmts()
+    descs = fetch_descriptions_batch(all_gs)
+    pairs = build_pairs(all_gs, descs)
+    print(f"\\nTotal pairs: {len(pairs)}")
+    split_and_save(pairs)