"""Modal GPU smoke test for Thousand Token Wood.

Proves the serving path works end to end before the hack weekend:
  - a GPU container spins up on Modal
  - a small model (<=32B) loads
  - it returns one in-character generation

Run:  python -m modal run modal_smoke_test.py
"""

import modal

MODEL = "Qwen/Qwen2.5-7B-Instruct"  # ~15GB in bf16, fits on a single L4 (24GB)

app = modal.App("ttw-smoke-test")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("transformers==4.46.0", "torch==2.5.1", "accelerate==1.1.1")
)


@app.function(gpu="L4", image=image, timeout=900)
def generate() -> str:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, torch_dtype=torch.bfloat16, device_map="cuda"
    )

    messages = [
        {
            "role": "system",
            "content": "You are Fenn, a sly fox trader in Thousand Token Wood.",
        },
        {
            "role": "user",
            "content": (
                "The price of acorns just crashed after a rumor that the harvest "
                "was poisoned. In one sentence, decide whether you buy or sell, "
                "and why."
            ),
        },
    ]
    text = tok.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tok(text, return_tensors="pt").to("cuda")
    out = model.generate(**inputs, max_new_tokens=80, do_sample=False)
    return tok.decode(
        out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True
    ).strip()


@app.local_entrypoint()
def main():
    print("\n=== Fenn the fox says ===")
    print(generate.remote())
    print("=========================\n")