python-gfm 0.1.2__tar.gz → 0.1.4__tar.gz

{python_gfm-0.1.2 → python_gfm-0.1.4}/PKG-INFO

@@ -1,16 +1,28 @@
 Metadata-Version: 2.4
 Name: python-gfm
-Version: 0.1.2
+Version: 0.1.4
 Summary: An AI copilot for graph data and models (Under active development).
 Author-email: BUAA SKLCCSE <your.email@example.com>
 License: Apache-2.0
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.20
 Requires-Dist: pyyaml>=6
 Provides-Extra: torch
-Requires-Dist: torch>=2.0; extra == "torch"
-Requires-Dist: torch-geometric>=2.3; extra == "torch"
+Requires-Dist: torch==2.8.0; extra == "torch"
+Requires-Dist: torch-geometric>=2.3.0; extra == "torch"
+Requires-Dist: torch-scatter==2.1.2; extra == "torch"
+Requires-Dist: torch-sparse==0.6.18; extra == "torch"
+Requires-Dist: torch-cluster==1.6.3; extra == "torch"
+Requires-Dist: torch-spline-conv==1.2.2; extra == "torch"
+Requires-Dist: scipy>=1.9.0; extra == "torch"
+Requires-Dist: transformers>=4.36.0; extra == "torch"
+Requires-Dist: joblib>=1.3.0; extra == "torch"
+Requires-Dist: scikit-learn>=1.3.0; extra == "torch"
+Requires-Dist: ftfy>=6.1.0; extra == "torch"
+Requires-Dist: regex>=2023.0.0; extra == "torch"
+Requires-Dist: accelerate>=0.26.0; extra == "torch"
+Requires-Dist: fschat>=0.2.36; extra == "torch"
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: ruff; extra == "dev"

python_gfm-0.1.4/pyproject.toml

@@ -0,0 +1,58 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "python-gfm"
+version = "0.1.4"
+authors = [
+  { name = "BUAA SKLCCSE", email = "your.email@example.com" },
+]
+description = "An AI copilot for graph data and models (Under active development)."
+readme = "README.md"
+requires-python = ">=3.12"
+license = { text = "Apache-2.0" }
+dependencies = [
+  "numpy>=1.20",
+  "pyyaml>=6",
+]
+
+# Optional ML / GFM stack (default target: Python 3.12, Linux x86_64, CUDA 12.8).
+# Torch and PyG extension versions are pinned to match:
+#   https://data.pyg.org/whl/torch-2.8.0+cu128.html
+#
+# Typical install:
+#   pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128
+#   pip install "python-gfm[torch]" -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
+#
+# The -f URL is required so pip picks the +pt28cu128 wheels for torch-scatter / torch-sparse /
+# torch-cluster / torch-spline-conv. CPU-only torch: use https://data.pyg.org/whl/torch-2.8.0+cpu.html
+# and install torch from PyPI or the PyTorch CPU index instead.
+[project.optional-dependencies]
+torch = [
+  "torch==2.8.0",
+  "torch-geometric>=2.3.0",
+  "torch-scatter==2.1.2",
+  "torch-sparse==0.6.18",
+  "torch-cluster==1.6.3",
+  "torch-spline-conv==1.2.2",
+  "scipy>=1.9.0",
+  "transformers>=4.36.0",
+  "joblib>=1.3.0",
+  "scikit-learn>=1.3.0",
+  "ftfy>=6.1.0",
+  "regex>=2023.0.0",
+  "accelerate>=0.26.0",
+  "fschat>=0.2.36",
+]
+dev = ["pytest", "ruff"]
+
+[project.scripts]
+pygfm = "pygfm.cli.run_yaml:main"
+gfm = "pygfm.cli.run_yaml:main"
+gfm-sa2gfm-pretrain = "pygfm.cli.sa2gfm:pretrain_main"
+gfm-sa2gfm-downstream = "pygfm.cli.sa2gfm:downstream_main"
+
+[tool.setuptools.packages.find]
+where = ["src", "scripts", "ckpts"]
+include = ["pygfm*"]

{python_gfm-0.1.2 → python_gfm-0.1.4}/python_gfm.egg-info/PKG-INFO

@@ -1,16 +1,28 @@
 Metadata-Version: 2.4
 Name: python-gfm
-Version: 0.1.2
+Version: 0.1.4
 Summary: An AI copilot for graph data and models (Under active development).
 Author-email: BUAA SKLCCSE <your.email@example.com>
 License: Apache-2.0
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.20
 Requires-Dist: pyyaml>=6
 Provides-Extra: torch
-Requires-Dist: torch>=2.0; extra == "torch"
-Requires-Dist: torch-geometric>=2.3; extra == "torch"
+Requires-Dist: torch==2.8.0; extra == "torch"
+Requires-Dist: torch-geometric>=2.3.0; extra == "torch"
+Requires-Dist: torch-scatter==2.1.2; extra == "torch"
+Requires-Dist: torch-sparse==0.6.18; extra == "torch"
+Requires-Dist: torch-cluster==1.6.3; extra == "torch"
+Requires-Dist: torch-spline-conv==1.2.2; extra == "torch"
+Requires-Dist: scipy>=1.9.0; extra == "torch"
+Requires-Dist: transformers>=4.36.0; extra == "torch"
+Requires-Dist: joblib>=1.3.0; extra == "torch"
+Requires-Dist: scikit-learn>=1.3.0; extra == "torch"
+Requires-Dist: ftfy>=6.1.0; extra == "torch"
+Requires-Dist: regex>=2023.0.0; extra == "torch"
+Requires-Dist: accelerate>=0.26.0; extra == "torch"
+Requires-Dist: fschat>=0.2.36; extra == "torch"
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: ruff; extra == "dev"

{python_gfm-0.1.2 → python_gfm-0.1.4}/python_gfm.egg-info/SOURCES.txt

@@ -45,7 +45,6 @@ src/pygfm/baseline_models/gcot/graphgpt/model/model_registry.py
 src/pygfm/baseline_models/gcot/graphgpt/model/monkey_patch_non_inplace.py
 src/pygfm/baseline_models/gcot/graphgpt/model/utils.py
 src/pygfm/baseline_models/gcot/graphgpt/model/graph_layers/__init__.py
-src/pygfm/baseline_models/gcot/graphgpt/model/graph_layers/bpe_simple_vocab_16e6.txt.gz
 src/pygfm/baseline_models/gcot/graphgpt/model/graph_layers/clip_graph.py
 src/pygfm/baseline_models/gcot/graphgpt/model/graph_layers/graph_transformer.py
 src/pygfm/baseline_models/gcot/graphgpt/model/graph_layers/mpnn.py
@@ -106,7 +105,6 @@ src/pygfm/baseline_models/graphgpt/model/model_registry.py
 src/pygfm/baseline_models/graphgpt/model/monkey_patch_non_inplace.py
 src/pygfm/baseline_models/graphgpt/model/utils.py
 src/pygfm/baseline_models/graphgpt/model/graph_layers/__init__.py
-src/pygfm/baseline_models/graphgpt/model/graph_layers/bpe_simple_vocab_16e6.txt.gz
 src/pygfm/baseline_models/graphgpt/model/graph_layers/clip_graph.py
 src/pygfm/baseline_models/graphgpt/model/graph_layers/graph_transformer.py
 src/pygfm/baseline_models/graphgpt/model/graph_layers/mpnn.py
@@ -363,7 +361,13 @@ src/pygfm/baseline_models/samgpt/downprompt.py
 src/pygfm/baseline_models/samgpt/downprompt_graph.py
 src/pygfm/baseline_models/samgpt/preprompt.py
 src/pygfm/cli/__init__.py
+src/pygfm/cli/baseline_registry.py
+src/pygfm/cli/mdgpt_stages.py
+src/pygfm/cli/run_yaml.py
 src/pygfm/cli/sa2gfm.py
+src/pygfm/cli/script_runner.py
+src/pygfm/cli/baselines/__init__.py
+src/pygfm/cli/baselines/stub_config.py
 src/pygfm/private/__init__.py
 src/pygfm/private/core/__init__.py
 src/pygfm/private/core/gnn_encoder.py

{python_gfm-0.1.2 → python_gfm-0.1.4}/python_gfm.egg-info/entry_points.txt

@@ -1,3 +1,5 @@
 [console_scripts]
+gfm = pygfm.cli.run_yaml:main
 gfm-sa2gfm-downstream = pygfm.cli.sa2gfm:downstream_main
 gfm-sa2gfm-pretrain = pygfm.cli.sa2gfm:pretrain_main
+pygfm = pygfm.cli.run_yaml:main

python_gfm-0.1.4/python_gfm.egg-info/requires.txt

@@ -0,0 +1,22 @@
+numpy>=1.20
+pyyaml>=6
+
+[dev]
+pytest
+ruff
+
+[torch]
+torch==2.8.0
+torch-geometric>=2.3.0
+torch-scatter==2.1.2
+torch-sparse==0.6.18
+torch-cluster==1.6.3
+torch-spline-conv==1.2.2
+scipy>=1.9.0
+transformers>=4.36.0
+joblib>=1.3.0
+scikit-learn>=1.3.0
+ftfy>=6.1.0
+regex>=2023.0.0
+accelerate>=0.26.0
+fschat>=0.2.36

{python_gfm-0.1.2 → python_gfm-0.1.4}/src/pygfm/__init__.py

@@ -4,7 +4,7 @@ pygfm: An AI copilot for graph data and models.
 This package provides tools and utilities for working with Graph Foundation Models.
 """
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 __author__ = "BUAA SKLCCSE"
 
 __all__ = ["__version__"]

python_gfm-0.1.4/src/pygfm/cli/baseline_registry.py

@@ -0,0 +1,103 @@
+"""Resolve ``(baseline, stage)`` → runner for ``pygfm -c config.yaml``."""
+
+from __future__ import annotations
+
+import importlib
+import pkgutil
+from typing import Any, Callable
+
+from pygfm.cli.baselines.stub_config import ALL_SCRIPT_PAIRS
+from pygfm.cli.script_runner import make_runner
+
+Runner = Callable[[dict[str, Any]], None]
+
+_STAGE_ALIASES: dict[str, str] = {
+    "ft": "finetune",
+    "finetune_node": "finetune",
+    "train": "pretrain",
+}
+
+
+def merge_config(cfg: dict[str, Any]) -> dict[str, Any]:
+    params = cfg.get("params")
+    base = {k: v for k, v in cfg.items() if k != "params"}
+    if isinstance(params, dict):
+        return {**params, **base}
+    return base
+
+
+def _discover_module_runners() -> dict[tuple[str, str], Runner]:
+    out: dict[tuple[str, str], Runner] = {}
+    import pygfm.cli.baselines as baselines_pkg
+
+    for info in pkgutil.iter_modules(baselines_pkg.__path__):
+        name = info.name
+        if name.startswith("_") or name == "stub_config":
+            continue
+        mod = importlib.import_module(f"pygfm.cli.baselines.{name}")
+        runners = getattr(mod, "RUNNERS", None)
+        if not isinstance(runners, dict):
+            continue
+        for stage, fn in runners.items():
+            if callable(fn):
+                out[(name, str(stage))] = fn
+    return out
+
+
+def _sa2gfm_pretrain(_cfg: dict[str, Any]) -> None:
+    from pygfm.cli.sa2gfm import pretrain_main
+
+    pretrain_main()
+
+
+def _sa2gfm_downstream(_cfg: dict[str, Any]) -> None:
+    from pygfm.cli.sa2gfm import downstream_main
+
+    downstream_main()
+
+
+def _build_registry() -> dict[tuple[str, str], Runner]:
+    reg: dict[tuple[str, str], Runner] = {}
+
+    for b, s in ALL_SCRIPT_PAIRS:
+        reg[(b, s)] = make_runner(b, s)
+
+    reg[("sa2gfm", "pretrain")] = _sa2gfm_pretrain
+    reg[("sa2gfm", "downstream")] = _sa2gfm_downstream
+
+    reg.update(_discover_module_runners())
+    return reg
+
+
+_REGISTRY: dict[tuple[str, str], Runner] | None = None
+
+
+def get_registry() -> dict[tuple[str, str], Runner]:
+    global _REGISTRY
+    if _REGISTRY is None:
+        _REGISTRY = _build_registry()
+    return _REGISTRY
+
+
+def run_from_yaml_dict(cfg: dict[str, Any]) -> None:
+    merged = merge_config(cfg)
+    baseline = merged.get("baseline")
+    stage = merged.get("stage")
+    if baseline is None or stage is None:
+        raise ValueError("YAML must set `baseline` and `stage`.")
+
+    b = str(baseline).strip().lower()
+    s = str(stage).strip().lower()
+    s = _STAGE_ALIASES.get(s, s)
+
+    fn = get_registry().get((b, s))
+    if fn is None:
+        raise ValueError(
+            f"No runner for baseline={b!r}, stage={s!r}. "
+            f"Known: {len(list_implemented())} pairs — see list_implemented()."
+        )
+    fn(merged)
+
+
+def list_implemented() -> list[str]:
+    return [f"{a}/{t}" for (a, t) in sorted(get_registry().keys())]

python_gfm-0.1.4/src/pygfm/cli/baselines/__init__.py

@@ -0,0 +1 @@
+"""Per-baseline YAML runners (pretrain / finetune / …) for ``pygfm`` CLI."""

python_gfm-0.1.4/src/pygfm/cli/baselines/stub_config.py

@@ -0,0 +1,40 @@
+"""
+Canonical list of ``(baseline, stage)`` pairs dispatched by ``script_runner`` (unless a
+real ``RUNNERS`` in ``pygfm.cli.baselines.<name>`` overrides).
+
+Used by ``baseline_registry`` to register ``make_runner(b, s)`` for every pair.
+"""
+
+from __future__ import annotations
+
+from pygfm.cli.script_runner import DEFAULT_STAGE_FILES, SCRIPT_OVERRIDES
+
+# Folders under scripts/ with the usual pretrain / finetune / finetune_graph layout.
+_STANDARD_BASELINES: tuple[str, ...] = (
+    "bridge",
+    "gcot",
+    "graver",
+    "graphkeeper",
+    "graphmore",
+    "graphprompt",
+    "hgprompt",
+    "mdgfm",
+    "mdgpt",
+    "rag_gfm",
+    "samgpt",
+)
+
+_STANDARD_STAGES: tuple[str, ...] = ("pretrain", "finetune", "finetune_graph")
+
+
+def _all_pairs() -> list[tuple[str, str]]:
+    pairs: set[tuple[str, str]] = set()
+    for b in _STANDARD_BASELINES:
+        for s in _STANDARD_STAGES:
+            pairs.add((b, s))
+    pairs.update(SCRIPT_OVERRIDES.keys())
+    pairs.add(("sa2gfm", "detect"))
+    return sorted(pairs)
+
+
+ALL_SCRIPT_PAIRS: list[tuple[str, str]] = _all_pairs()

python_gfm-0.1.4/src/pygfm/cli/mdgpt_stages.py

@@ -0,0 +1,270 @@
+"""MDGPT pretrain / finetune runners driven by YAML config (no ``scripts/`` layout)."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch_geometric.data import Batch, Data
+
+from pygfm.baseline_models.mdgpt import DownPromptModel, PrePromptModel
+from pygfm.private.utlis.domain_alignment import DomainAlignment
+from pygfm.public.utils import early_stopping, load_all_datasets, set_seed
+from pygfm.public.utils.loss_func import sample_negative_pairs
+from pygfm.public.utils.runtime import load_single_graph_dataset
+
+
+def _data_root(cfg: dict[str, Any]) -> str:
+    return str(cfg.get("data_root") or os.environ.get("GFM_DATA_ROOT", "datasets/mdgpt"))
+
+
+def run_mdgpt_pretrain(cfg: dict[str, Any]) -> None:
+    p = cfg.get("pretrain") or {}
+    if not p:
+        raise ValueError("YAML: missing `pretrain:` block for stage=pretrain")
+
+    seed = int(cfg.get("seed", 42))
+    data_root = _data_root(cfg)
+    target = str(p.get("target", "Cora"))
+
+    unify_dim = int(p.get("unify_dim", 50))
+    hidden_dim = int(p.get("hidden_dim", 256))
+    num_layers = int(p.get("num_layers", 3))
+    num_neg = int(p.get("num_neg", 50))
+    lr = float(p.get("lr", 1e-5))
+    patience = int(p.get("patience", 50))
+    max_epochs = int(p.get("max_epochs", 200))
+    log_interval = int(p.get("log_interval", 50))
+    prompt_mode = str(p.get("prompt_mode", "mul"))
+    temperature = float(p.get("temperature", 1.0))
+
+    allow_pyg = bool(p.get("allow_pyg_download", False))
+
+    save_dir = p.get("save_dir")
+    if save_dir:
+        save_dir = str(save_dir)
+    else:
+        save_dir = os.path.join("ckpts", "mdgpt", target.lower())
+    save_name = str(p.get("save_name") or f"preprompt_{target.lower()}.pth")
+
+    set_seed(seed)
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    all_raw = load_all_datasets(data_root, allow_pyg_download=allow_pyg)
+    if not all_raw:
+        raise RuntimeError(
+            f"No datasets under {data_root!r}. Add flat *.pt files or set allow_pyg_download."
+        )
+
+    sources = [d for d in all_raw if d["name"].casefold() != target.casefold()]
+    ordered_names = [s["name"] for s in sources]
+    if not sources:
+        raise RuntimeError(
+            f"No source domains after excluding target={target!r}. Found: {[d['name'] for d in all_raw]}"
+        )
+    if not any(d["name"].casefold() == target.casefold() for d in all_raw):
+        print(
+            f"!! Warning: target {target!r} not in {data_root!r}. "
+            "Leave-one-out pretrain still runs; finetune needs target .pt later."
+        )
+    print(f"Target: {target} | sources: {ordered_names}")
+
+    source_list: list[Data] = []
+    aligners = []
+    for idx, s in enumerate(sources):
+        g0 = s["ds"][0]
+        feat = g0.x.numpy()
+        aligner = DomainAlignment(n_components=unify_dim)
+        aligner.fit(feat)
+        aligners.append(aligner)
+        aligned = torch.from_numpy(aligner.transform(feat)).float()
+        d = Data(x=aligned, edge_index=g0.edge_index, y=g0.y)
+        d.domain_id = torch.full((aligned.size(0),), idx, dtype=torch.long)
+        source_list.append(d)
+
+    big_batch = Batch.from_data_list(source_list).to(device)
+    tuples = sample_negative_pairs(
+        big_batch.edge_index,
+        big_batch.num_nodes,
+        num_neg=num_neg,
+        seed=seed,
+    ).to(device)
+
+    model = PrePromptModel(
+        input_dim=unify_dim,
+        hidden_dim=hidden_dim,
+        num_domains=len(sources),
+        num_layers=num_layers,
+        prompt_mode=prompt_mode,  # type: ignore[arg-type]
+        temperature=temperature,
+        device=device,
+    )
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=0.0)
+
+    best_loss = 1e9
+    cnt_wait = 0
+    for epoch in range(max_epochs):
+        model.train()
+        optimizer.zero_grad()
+        loss = model(big_batch.x, big_batch.edge_index, big_batch.batch, tuples)
+        loss.backward()
+        optimizer.step()
+        should_stop, best_loss, cnt_wait = early_stopping(loss.item(), best_loss, cnt_wait, patience)
+        if epoch % log_interval == 0:
+            print(f"epoch={epoch} loss={loss.item():.4f}")
+        if should_stop:
+            print("early stopping")
+            break
+
+    os.makedirs(save_dir, exist_ok=True)
+    ckpt_path = os.path.join(save_dir, save_name)
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "unify_dim": unify_dim,
+            "hidden_dim": hidden_dim,
+            "num_domains": len(sources),
+            "ordered_names": ordered_names,
+            "prompt_mode": prompt_mode,
+            "target": target,
+        },
+        ckpt_path,
+    )
+    print("saved:", ckpt_path)
+
+    try:
+        import joblib
+
+        aligner_path = os.path.join(save_dir, "aligners.pkl")
+        joblib.dump({"aligners": aligners, "ordered_names": ordered_names}, aligner_path)
+        print("saved aligners:", aligner_path)
+    except Exception:
+        pass
+
+
+def run_mdgpt_finetune(cfg: dict[str, Any]) -> None:
+    f = cfg.get("finetune") or {}
+    if not f:
+        raise ValueError("YAML: missing `finetune:` block for stage=finetune")
+
+    seed = int(cfg.get("seed", 42))
+    data_root = _data_root(cfg)
+
+    dataset = str(f.get("dataset", "Cora"))
+    k_shot = int(f.get("k_shot", 1))
+    split_id = int(f.get("split_id", 0))
+    ckpt_path = f.get("ckpt")
+    if not ckpt_path:
+        raise ValueError("YAML finetune.ckpt: path to preprompt *.pth required")
+    ckpt_path = str(ckpt_path)
+
+    downstream_root = str(f.get("downstream_root", "downstream_data/mdgpt"))
+    splits_path = f.get("splits_path")
+    if splits_path:
+        splits_path = str(splits_path)
+    else:
+        splits_path = os.path.join(downstream_root, dataset, f"{k_shot}shot", "splits.pt")
+
+    lr = float(f.get("lr", 1e-3))
+    patience = int(f.get("patience", 50))
+    max_steps = int(f.get("max_steps", 400))
+    test_reserve = int(f.get("test_reserve", 1000))
+    allow_pyg = bool(f.get("allow_pyg_download", False))
+
+    set_seed(seed)
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    ckpt = torch.load(ckpt_path, map_location=device)
+    unify_dim = int(ckpt["unify_dim"])
+    hidden_dim = int(ckpt["hidden_dim"])
+    prompt_mode = str(ckpt["prompt_mode"])
+
+    preprompt = PrePromptModel(
+        input_dim=unify_dim,
+        hidden_dim=hidden_dim,
+        num_domains=int(ckpt.get("num_domains", 4)),
+        num_layers=int(f.get("num_layers", 3)),
+        prompt_mode=prompt_mode,  # type: ignore[arg-type]
+        temperature=float(f.get("temperature", 1.0)),
+        device=device,
+    )
+    preprompt.load_state_dict(ckpt["model"], strict=False)
+    preprompt.eval()
+
+    data, num_classes = load_single_graph_dataset(data_root, dataset, allow_pyg_download=allow_pyg)
+    edge_index = data.edge_index.to(device)
+    y = data.y.to(device)
+
+    x_raw = data.x.cpu().numpy()
+    aligner = DomainAlignment(n_components=unify_dim)
+    aligner.fit(x_raw)
+    x = torch.from_numpy(aligner.transform(x_raw)).float().to(device)
+
+    down_data = torch.load(splits_path, map_location="cpu")
+    splits = down_data["splits"]
+    if not (0 <= split_id < len(splits)):
+        raise IndexError(f"split_id {split_id} out of range 0..{len(splits)-1}")
+    split = splits[split_id]
+    support_idx = torch.tensor(split["indices"], dtype=torch.long, device=device)
+    support_labels = torch.tensor(split["labels"], dtype=torch.long, device=device)
+
+    test_start = max(0, len(y) - test_reserve)
+    test_idx = torch.arange(test_start, len(y), device=device)
+    test_labels = y[test_idx]
+
+    down = DownPromptModel(
+        gcn=preprompt.gcn,
+        input_dim=unify_dim,
+        hidden_dim=hidden_dim,
+        num_classes=num_classes,
+        prompt_mode=prompt_mode,
+        device=device,
+    )
+    try:
+        down.prefeature.load_state_dict(preprompt.pretexts[0].state_dict(), strict=False)
+    except Exception:
+        pass
+
+    opt = torch.optim.Adam(down.prefeature.parameters(), lr=lr)
+    best_loss = 1e9
+    wait = 0
+    for step in range(max_steps):
+        down.train()
+        opt.zero_grad()
+        logits = down(
+            x,
+            edge_index,
+            support_idx=support_idx,
+            support_labels=support_labels,
+            query_idx=support_idx,
+            train=True,
+        )
+        loss = F.cross_entropy(logits, support_labels)
+        loss.backward()
+        opt.step()
+        if loss.item() < best_loss:
+            best_loss = loss.item()
+            wait = 0
+        else:
+            wait += 1
+        if step % 50 == 0:
+            print(f"step={step} loss={loss.item():.4f}")
+        if wait >= patience:
+            print(f"early stopping step={step} best_loss={best_loss:.4f}")
+            break
+
+    down.eval()
+    with torch.inference_mode():
+        logits = down(
+            x,
+            edge_index,
+            support_idx=support_idx,
+            support_labels=support_labels,
+            query_idx=test_idx,
+            train=False,
+        )
+        preds = logits.argmax(dim=1)
+        acc = (preds == test_labels).float().mean().item()
+    print(f"[{dataset}] {k_shot}-shot split {split_id} test acc: {acc:.4f}")

python_gfm-0.1.4/src/pygfm/cli/run_yaml.py

@@ -0,0 +1,46 @@
+"""Unified CLI: run experiments from YAML for any registered ``baseline`` + ``stage``."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from pygfm.cli.baseline_registry import list_implemented, run_from_yaml_dict
+
+
+def main(argv: list[str] | None = None) -> None:
+    p = argparse.ArgumentParser(
+        description=(
+            "Run GFM from YAML with `baseline` and `stage` (e.g. mdgpt + pretrain). "
+            "Registered pairs include all stub stages from stub_config plus real RUNNERS in "
+            "pygfm.cli.baselines.<name> — see list_implemented() in baseline_registry."
+        ),
+    )
+    p.add_argument(
+        "-c",
+        "--config",
+        "--yaml-config",
+        dest="config",
+        type=str,
+        required=True,
+        help="Path to YAML config file",
+    )
+    args = p.parse_args(argv)
+
+    path = Path(args.config)
+    if not path.is_file():
+        raise SystemExit(f"Config not found: {path}")
+
+    raw = path.read_text(encoding="utf-8")
+    cfg: dict[str, Any] = yaml.safe_load(raw) or {}
+    if not isinstance(cfg, dict):
+        raise SystemExit("YAML root must be a mapping (object)")
+
+    run_from_yaml_dict(cfg)
+
+
+if __name__ == "__main__":
+    main()