| """Small MPS compatibility probe. Concurrent CPU evaluation means timings are provisional.""" |
|
|
| import json |
| import os |
| import time |
| from pathlib import Path |
|
|
| os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "0" |
| import torch |
| from gliner2.classification import ClassificationSchema |
|
|
| from sealcell.model import LocalModel, validate_head |
|
|
| capabilities = { |
| "torch": torch.__version__, |
| "mps_built": torch.backends.mps.is_built(), |
| "mps_available": torch.backends.mps.is_available(), |
| "fallback_enabled": False, |
| } |
| print(json.dumps(capabilities), flush=True) |
| if not capabilities["mps_available"]: |
| raise SystemExit("MPS unavailable") |
| model = LocalModel() |
| source = json.loads(Path("data/fixtures/playground-probes.json").read_text()) |
| schema = ClassificationSchema() |
| for name, labels in source["schema"].items(): |
| schema.single(name, labels) |
| model.classifier.to(device="mps", dtype=torch.float32).eval() |
| records = [] |
| for case in source["cases"]: |
| torch.mps.synchronize() |
| start = time.perf_counter() |
| try: |
| result = model.classifier.classify(case["text"], schema) |
| torch.mps.synchronize() |
| elapsed = (time.perf_counter() - start) * 1000 |
| heads = { |
| name: { |
| "selected_label": result.value(name), |
| "probabilities": dict(result.probabilities(name)), |
| } |
| for name in source["schema"] |
| } |
| for name, head in heads.items(): |
| validate_head(head["selected_label"], head["probabilities"], source["schema"][name]) |
| record = { |
| "fixture_id": case["id"], |
| "device": "mps", |
| "dtype": "float32", |
| "heads": heads, |
| "total_ms": elapsed, |
| "error": None, |
| } |
| except Exception as exc: |
| record = { |
| "fixture_id": case["id"], |
| "device": "mps", |
| "error": f"{type(exc).__name__}: {exc}", |
| } |
| records.append(record) |
| print(case["id"], record.get("total_ms"), record["error"], flush=True) |
| Path("reports/mps-compatibility.json").write_text( |
| json.dumps( |
| { |
| "capabilities": capabilities, |
| "model_revision": model.metadata["revision"], |
| "timing_context": "provisional; CPU evaluation running concurrently; first shape includes compilation", |
| "records": records, |
| }, |
| indent=2, |
| ) |
| + "\n" |
| ) |
|
|
