solver_tf.py

#!/usr/bin/env python3
"""transformers solver subprocess (fallback path): fp16 model, batched generate.

Reads IOL_TEST_CSV / IOL_BUDGET_S from env; writes submission.csv incrementally.
"""
import csv
import os
import time

START = time.time()
BUDGET = float(os.environ.get("IOL_BUDGET_S", 1200))
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")


def log(msg):
    print(f"[tf-solver +{time.time()-START:5.1f}s] {msg}", flush=True)


def main():
    with open(TEST_CSV, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    log(f"{len(rows)} rows")
    import transformers
    log(f"transformers=={transformers.__version__}")

    model_dir = "q4b" if os.path.isdir("q4b") else "."
    from probe_pipeline import solve_all_transformers
    solve_all_transformers(model_dir, rows, "submission.csv", START, BUDGET - 60, log)


if __name__ == "__main__":
    main()