env.py

import os
import torch

if torch.backends.mps.is_available():
    device = torch.device("mps")
elif torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
print(f"[env] Using device: {device}")

MASTER_FP32 = "bert_fp32.pt"


def prepare_model(model, mode="auto"):
    if mode == "auto":
        if device.type == "cpu":
            mode = "int8"
        else:
            mode = "fp16"

    if mode == "fp32":
        print("[env] Using FP32 model (baseline).")
        return model.float(), "fp32"

    if mode == "fp16":
        if device.type in ["cuda", "mps"]:
            print("[env] Using FP16 model on GPU/MPS.")
            return model.half(), "fp16"
        else:
            print("[env][warn] FP16 requested but no GPU/MPS available. Falling back to FP32.")
            return model.float(), "fp32"

    if mode == "int8":
        if device.type == "cpu":
            print("[env] Using INT8 dynamic quantization on CPU.")
            return torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8), "int8"
        else:
            print("[env][warn] INT8 requested but not supported on GPU/MPS. Falling back to FP32.")
            return model.float(), "fp32"

    raise ValueError(f"[env] Unknown mode: {mode}")


def restore_model(target="fp32"):
    if target == "fp32" and os.path.exists(MASTER_FP32):
        print("[env] Reloading clean FP32 checkpoint.")
        state_dict = torch.load(MASTER_FP32, map_location="cpu")
        return state_dict, True
    else:
        print(f"[env][warn] {target} requested but only quantized/degraded weights available.")
        return None, False