generate_inkling.py

#!/usr/bin/env python3
"""Multimodal generation with the inkling (tml) model in plain Transformers.

Runs one prompt, optionally with attached image(s) and/or audio clip(s), and prints the
model's answer. The model is sharded across all visible GPUs with `tp_plan="auto"`, so
launch it under `torchrun` (see submit_inkling_generate.sbatch) for a multi-GPU / multi-node
checkpoint.

    torchrun --nnodes=1 --nproc-per-node=8 generate_inkling.py \
        --model-path /path/to/tml-model-share \
        --image cat.png --prompt "What is in this image?"

    torchrun ... generate_inkling.py --audio clip.wav --prompt "Transcribe this."
    torchrun ... generate_inkling.py --prompt "Explain tensor parallelism."
"""

from __future__ import annotations

import argparse
import io
import os
import urllib.request
import wave

import numpy as np


os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import torch

from transformers import AutoProcessor, InklingConfig, InklingForConditionalGeneration


def load_image(source: str):
    from PIL import Image

    if source.startswith(("http://", "https://")):
        with urllib.request.urlopen(source, timeout=60) as response:
            source = io.BytesIO(response.read())
    return Image.open(source).convert("RGB")


def load_wav(path: str, expected_rate: int) -> np.ndarray:
    """Decode a mono float32 waveform from a PCM .wav with the stdlib (bypasses torchcodec)."""
    with wave.open(path, "rb") as f:
        rate, n_channels, sample_width = f.getframerate(), f.getnchannels(), f.getsampwidth()
        frames = f.readframes(f.getnframes())
    if rate != expected_rate:
        raise SystemExit(f"{path}: sample rate {rate} != {expected_rate}; resample to 16 kHz first")
    if sample_width == 2:
        samples = np.frombuffer(frames, dtype=np.int16).
        astype(np.float32) / 32768.0
    elif sample_width == 4:
        samples = np.frombuffer(frames, dtype=np.int32).astype(np.float32) / 2147483648.0
    else:
        raise SystemExit(f"{path}: unsupported sample width {sample_width} bytes (use 16/32-bit PCM)")
    if n_channels > 1:
        samples = samples.reshape(-1, n_channels).mean(axis=1)
    return samples


def build_messages(prompt: str, images: list[str], audios: list[str], sampling_rate: int, effort: float | None):
    """One user message whose content is the attached media followed by the text question.

    An optional scalar reasoning effort (0.0-1.0) is rendered as a leading system message.
    """
    content: list[dict] = []
    for image in images:
        content.append({"type": "image", "image": load_image(image)})
    for audio in audios:
        content.append({"type": "audio", "audio": load_wav(audio, sampling_rate)})
    content.append({"type": "text", "text": prompt})

    messages = []
    if effort is not None:
        messages.append({"role": "system", "content": [{"type": "text", "text": f"Thinking effort level: {effort}"}]})
    messages.append({"role": "user", "content": content})
    return messages


def main() -> None:
    parser = argparse.ArgumentParser(description="Multimodal generation with the inkling (tml) model.")
    parser.add_argument("--model-path", default=os.environ.get("MODEL_PATH", "/fsx/pedro/tm/models/tml-model-share"))
    parser.add_argument("--processor-path", default=os.environ.get("PROCESSOR_PATH"), help="Defaults to --model-path.")
    parser.add_argument("--prompt", default="Explain tensor parallelism in simple terms.")
    parser.add_argument("--image", action="append", default=[], help="Image file path or http(s) URL (repeatable).")
    parser.add_argument("--audio", action="append", default=[], help="16 kHz mono PCM .wav path (repeatable).")
    parser.add_argument("--max-new-tokens", type=int, default=int(os.environ.get("MAX_NEW_TOKENS", "512")))
    parser.add_argument("--thinking-effort", type=float, default=None, help="Reasoning effort 0.0-1.0.")
    parser.add_argument("--do-sample", action="store_true")
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-p", type=float, default=1.0)
    # torchrun may append launcher args after the script path; ignore anything unrecognised.
    args, _ = parser.parse_known_args()

    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    processor = AutoProcessor.from_pretrained(args.processor_path or args.model_path)
    config = InklingConfig.from_pretrained(args.model_path)
    model = InklingForConditionalGeneration.from_pretrained(
        args.model_path,
        config=config,
        dtype=torch.bfloat16,
        tp_plan="auto",
    )
    model.eval()

    messages = build_messages(
        args.prompt, args.image, args.audio, processor.feature_extractor.sampling_rate, args.thinking_effort
    )
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        add_generation_prompt=True,
    ).to(device)
    if "pixel_values" in inputs:
        # the vision tower runs in the model dtype
        inputs["pixel_values"] = inputs["pixel_values"].to(dtype=torch.bfloat16)

    generate_kwargs = {
        "max_new_tokens": args.max_new_tokens,
        "do_sample": args.do_sample,
        "pad_token_id": config.eos_token_id,
    }
    if args.do_sample:
        generate_kwargs["temperature"] = args.temperature
        generate_kwargs["top_p"] = args.top_p

    with torch.no_grad():
        generated = model.generate(**inputs, **generate_kwargs)

    if rank == 0:
        new_tokens = generated[0, inputs["input_ids"].shape[1] :]
        answer = processor.tokenizer.decode(new_tokens, skip_special_tokens=False)
        print(
            f"[prompt {inputs['input_ids'].shape[1]} tok | images {len(args.image)} | audios {len(args.audio)} "
            f"| generated {new_tokens.numel()} tok]",
            flush=True,
        )
        print(answer, flush=True)


if __name__ == "__main__":
    main()