| """Run frozen MatrixFormer-Impute or MatrixFormer+ on a numerical matrix.""" |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| from huggingface_hub import hf_hub_download |
|
|
| from tabimpute.routing import TwoExpertRouter |
| from tabimpute.streamed_inference import StreamedImputePFN |
|
|
|
|
| REPO_ID = "Tabimpute/TabImpute" |
|
|
|
|
| def artifact_path(filename, local_dir, revision): |
| if local_dir is not None: |
| path = local_dir / filename |
| if not path.is_file(): |
| raise FileNotFoundError(f"Missing local artifact: {path}") |
| return path |
| return Path(hf_hub_download(REPO_ID, filename, revision=revision)) |
|
|
|
|
| def verify_sha256(path, expected): |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): |
| digest.update(chunk) |
| if digest.hexdigest() != expected: |
| raise ValueError(f"SHA-256 mismatch: {path}") |
|
|
|
|
| def load_expert(name, manifest, args): |
| spec = manifest["experts"][name] |
| config_path = artifact_path(spec["config"], args.local_dir, args.revision) |
| checkpoint_path = artifact_path(spec["checkpoint"], args.local_dir, args.revision) |
| verify_sha256(config_path, spec["config_sha256"]) |
| verify_sha256(checkpoint_path, spec["checkpoint_sha256"]) |
| config = json.loads(config_path.read_text(encoding="utf-8")) |
| return StreamedImputePFN( |
| checkpoint_path=str(checkpoint_path), |
| json_config={"config": config["model"]}, |
| device=args.device, |
| prediction_stat="mean", |
| precision_mode=args.precision, |
| max_num_rows=args.max_rows or None, |
| max_num_cols=args.max_cols or None, |
| ) |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--model", choices=("impute", "plus"), default="impute") |
| parser.add_argument("--input", type=Path, help="2D numeric .npy array; missing entries are NaN") |
| parser.add_argument("--output", type=Path, help="Completed .npy array") |
| parser.add_argument("--device", default="cpu", help="For example: cpu, cuda, cuda:0") |
| parser.add_argument("--precision", choices=("amp_bfloat16", "float32"), default="amp_bfloat16") |
| parser.add_argument("--max-rows", type=int, default=2000, help="Context rows; 0 uses the full axis") |
| parser.add_argument("--max-cols", type=int, default=1000, help="Context columns; 0 uses the full axis") |
| parser.add_argument("--local-dir", type=Path, help="Offline bundle root containing matrixformer/") |
| parser.add_argument("--revision", default="main", help="Hugging Face revision; use a commit hash to pin it") |
| args = parser.parse_args() |
| if args.max_rows < 0 or args.max_cols < 0: |
| parser.error("Context limits must be nonnegative") |
| if args.output is not None and args.output.suffix != ".npy": |
| parser.error("--output must have a .npy suffix") |
| if args.input is not None and args.output is not None and args.input.resolve() == args.output.resolve(): |
| parser.error("--output must differ from --input") |
| return args |
|
|
|
|
| def main(): |
| args = parse_args() |
| if args.input is None: |
| generator = np.random.default_rng(42) |
| matrix = generator.normal(size=(12, 3)) @ generator.normal(size=(3, 9)) |
| missing = generator.choice(matrix.size, size=30, replace=False) |
| matrix.flat[missing] = np.nan |
| else: |
| matrix = np.load(args.input, allow_pickle=False) |
| if not np.issubdtype(matrix.dtype, np.number) or np.iscomplexobj(matrix): |
| raise ValueError("Input must contain real numbers, with NaN marking missing entries") |
| matrix = np.asarray(matrix, dtype=np.float64) |
| if matrix.ndim != 2 or min(matrix.shape) == 0 or np.isinf(matrix).any(): |
| raise ValueError("Expected a nonempty 2D matrix with finite observations or NaN") |
| if not np.isfinite(matrix).any(): |
| raise ValueError("At least one observed entry is required") |
| manifest_path = artifact_path("matrixformer/artifacts.json", args.local_dir, args.revision) |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) |
| general = load_expert("impute", manifest, args) |
| if args.model == "plus": |
| specialist = load_expert("recsys", manifest, args) |
| imputer = TwoExpertRouter( |
| general, specialist, policy="observed_blend", normalization="column" |
| ) |
| else: |
| imputer = general |
| completed = imputer.impute(matrix.copy()) |
| if not np.isfinite(completed).all(): |
| raise RuntimeError("Model returned nonfinite predictions") |
| observed = np.isfinite(matrix) |
| if not np.array_equal(completed[observed], matrix[observed]): |
| raise RuntimeError("Model changed observed entries") |
| report = { |
| "model": args.model, |
| "device": args.device, |
| "precision": args.precision, |
| "shape": list(matrix.shape), |
| "missing_entries": int(np.isnan(matrix).sum()), |
| "all_predictions_finite": True, |
| "observed_entries_preserved": True, |
| } |
| if args.model == "plus": |
| report["blend_decision"] = imputer.last_result.decisions["observed"] |
| if args.output is not None: |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| np.save(args.output, completed, allow_pickle=False) |
| report["output"] = str(args.output) |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
