Function test_inference

modal_setup/dev_app.py:56–87

Test TIDE inference on a model.

(
    model_name: str = "meta-llama/Llama-3.1-8B-Instruct",
    prompt: str = "Explain quantum computing in simple terms:",
    max_new_tokens: int = 256,
)

    timeout=3600,
)
def test_inference(
    model_name: str = "meta-llama/Llama-3.1-8B-Instruct",
    prompt: str = "Explain quantum computing in simple terms:",
    max_new_tokens: int = 256,
):
    """Test TIDE inference on a model."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from TIDE.config import TIDEConfig
    from TIDE.runtime import TIDERuntime

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto",
        cache_dir="/root/models",
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/root/models")

    safe_name = model_name.replace("/", "_")
    router_path = f"/root/routers/{safe_name}_router.pt"

    runtime = TIDERuntime(model, router_path, config=TIDEConfig())

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = runtime.generate(inputs.input_ids, max_new_tokens=max_new_tokens, temperature=0)
    text = tokenizer.decode(output[0], skip_special_tokens=True)

    return {
        "output": text,
        "stats": runtime.last_stats.summary() if runtime.last_stats else "N/A",
    }