Vision

import io

from datasets import load_dataset
from PIL import Image
from transformers import AutoImageProcessor

from baseten.loops import (
    AdamParams,
    Datum,
    ImageChunk,
    ModelInput,
    ServiceClient,
    TensorData,
)

BASE_MODEL = "Qwen/Qwen3.5-0.8B"

service_client = ServiceClient()
training_client = service_client.create_lora_training_client(
    base_model=BASE_MODEL,
    rank=16,
)
tokenizer = training_client.get_tokenizer()
image_processor = AutoImageProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)

def compute_expected_tokens(image: Image.Image) -> int:
    result = image_processor(images=[image])
    grid = result["image_grid_thw"][0]
    merge = image_processor.merge_size
    return int(grid[0] * grid[1] * grid[2]) // (merge * merge)


def to_image_chunk(image: Image.Image, expected_tokens: int) -> ImageChunk:
    buf = io.BytesIO()
    image.convert("RGB").save(buf, format="JPEG")
    return ImageChunk(
        data=buf.getvalue(),
        format="jpeg",
        expected_tokens=expected_tokens,
    )


def to_datum(example, tokenizer):
    image = example["images"][0]
    problem = example["problem"]
    answer = example["answer"]

    expected_tokens = compute_expected_tokens(image)
    img_chunk = to_image_chunk(image, expected_tokens)

    # geometry3k wraps the prompt in an <image> placeholder; strip it.
    prompt = problem.replace("<image>", "")
    p = tokenizer.encode(prompt, add_special_tokens=False)
    a = tokenizer.encode(answer, add_special_tokens=False)

    # The trainer doesn't shift labels: drop the last input token and make
    # each target the token that follows its position.
    text_chunk = ModelInput.from_ints(p + a[:-1]).chunks[0]
    targets = [-100] * (expected_tokens + len(p) - 1) + list(a)

    return Datum(
        model_input=ModelInput(chunks=[img_chunk, text_chunk]),
        loss_fn_inputs={
            "target_tokens": TensorData(
                data=targets, dtype="int64", shape=[len(targets)]
            )
        },
    )

BATCH_SIZE = 4

dataset = load_dataset("hiyouga/geometry3k", split="train[:4]")
data = [to_datum(ex, tokenizer) for ex in dataset]
print(f"prepared {len(data)} examples")

for step, start in enumerate(range(0, len(data), BATCH_SIZE), 1):
    batch = data[start : start + BATCH_SIZE]
    fb = training_client.forward_backward(data=batch).result(timeout=600.0)
    training_client.optim_step(
        AdamParams(learning_rate=4e-5)
    ).result(timeout=600.0)
    print(f"step {step} loss={fb.loss:.4f}")

save_resp = training_client.save_weights_for_sampler(name="vision-step-1").result(timeout=600.0)
print(f"saved checkpoint at {save_resp.path}")

import base64
import io
import json

from datasets import load_dataset
from PIL import Image

ds = load_dataset("hiyouga/geometry3k", split="train[:1]")
img = ds[0]["images"][0]
buf = io.BytesIO()
img.convert("RGB").save(buf, format="JPEG")
b64 = base64.b64encode(buf.getvalue()).decode()
problem = ds[0]["problem"].replace("<image>", "")

payload = {
    "model": "vision-step-1",
    "messages": [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
                },
                {"type": "text", "text": problem},
            ],
        }
    ],
    "max_tokens": 64,
    "chat_template_kwargs": {"enable_thinking": False},
}

with open("request.json", "w") as f:
    json.dump(payload, f)