#!/usr/bin/env python3
"""Make mlx-vlm's trainer accept an exact local JSONL file."""
from pathlib import Path
import runpy
import sys
import os

import datasets

_load_dataset = datasets.load_dataset


def load_dataset(path, name=None, split=None, **kwargs):
    candidate = Path(path)
    if candidate.exists() or str(path).endswith((".json", ".jsonl")):
        return _load_dataset("json", data_files=str(candidate), split=split or "train")
    return _load_dataset(path, name, split=split, **kwargs)


datasets.load_dataset = load_dataset

# Qwen3-VL uses a dynamic-resolution processor.  mlx-vlm's generic
# image-resize-shape option is intentionally ignored by that native path, so
# expose a safe pixel budget for local 16 GB Apple Silicon training.
try:
    from mlx_vlm import utils as _mlx_utils

    def _bounded_process_inputs(processor, prompts, images, audio, add_special_tokens=False, return_tensors="mlx", **kwargs):
        max_pixels = os.environ.get("ABJAD_MAX_PIXELS")
        if max_pixels:
            # Qwen3-VL exposes these through **kwargs, so the generic helper
            # cannot forward them by signature inspection.
            kwargs["min_pixels"] = 56 * 56
            kwargs["max_pixels"] = int(max_pixels)
        process_method = getattr(processor, "process", processor)
        call = {
            "text": prompts,
            "images": images,
            "padding": True,
            "return_tensors": return_tensors,
            "add_special_tokens": add_special_tokens,
            **kwargs,
        }
        return process_method(**call)

    _mlx_utils.process_inputs_with_fallback = _bounded_process_inputs
except Exception:
    pass

sys.argv[0] = "mlx_vlm.lora"
runpy.run_module("mlx_vlm.lora", run_name="__main__")