File size: 5,425 Bytes
7275749 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """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()
|