benchmarks/modal/bench_lora_test.py

"""
LoRA Adapter Test — evaluate the just-trained LoRA on the production
PlannerEngine prompts to see if it beats Ministral-8B-Instruct-2410 (95%)
and Gemma-3-4B (95%).

LoRA adapter: pranaysuyash/shopstack-parser-lora-qwen2.5-1.5b
Base model: Qwen/Qwen2.5-1.5B-Instruct

Run:
    cd /Users/pranay/Projects/shopstack
    unset MODAL_TOKEN_ID MODAL_TOKEN_SECRET
    modal run benchmarks/modal/bench_lora_test.py
"""

from __future__ import annotations

import json
import os
import re
import time
from pathlib import Path

import modal

app = modal.App("shopstack-lora-test")

GPU_CONFIG = "A10G"  # 1.5B model is small, A10G is enough

LORA_REPO = "pranaysuyash/shopstack-parser-lora-qwen2.5-1.5b"
BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"

# Same 20 prompts from the planner bench
TEST_PROMPTS = [
    ("add_milk",            "I bought 2 liters of milk. Record it in my inventory in the fridge."),
    ("find_onion",          "Do I have any onions at home?"),
    ("consume_rice",        "I used 0.5 kg of basmati rice from my pantry."),
    ("shopping_vegetables", "Create a shopping list for tomatoes, onions, and potatoes. I need 1 kg of each."),
    ("compare_eggs",        "I'm at the store and see eggs for $3.99. Should I buy them?"),
    ("price_tomato",        "I saw tomatoes at $2.49 per kg at Dmart. Record this price."),
    ("use_soon_check",      "What items in my fridge need to be used soon?"),
    ("buy_suggestions",     "What should I buy next time I go shopping?"),
    ("move_sugar",          "I moved the sugar from the pantry to the kitchen counter."),
    ("multi_step",          "I bought 3 kg of apples and 1 kg of carrots. Record both in the fridge, then check if I need to buy onions."),
    ("hinglish_add",        "tamatar aadha kilo add karo fridge mein"),
    ("hinglish_consume",    "maine 200 gm dahi use kiya subah"),
    ("ambiguous_qty",       "I bought some sugar. Add it to the pantry."),
    ("correction",          "Wait, that wasn't sugar, it was salt. Change the last item to salt instead."),
    ("empty_inventory",     "List everything I have at home right now."),
    ("find_multiple",       "Do I have any of these: onions, garlic, ginger, or turmeric?"),
    ("add_bulk",            "I went to the wholesale market and bought: 5kg rice, 3kg atta, 2L cooking oil, 1kg sugar, 500g tea. Add all to pantry."),
    ("price_history_check", "What's the average price of tomatoes over the last month?"),
    ("use_soon_pantry",     "Check what's expiring in my pantry this week."),
    ("confirm_action",      "I want to add 10 kg of rice but I'm not sure if I should. Ask me a question first."),
]

EXPECTED_TOOLS = {
    "add_milk":            [{"tool": "add_inventory_item"}],
    "find_onion":          [{"tool": "search_inventory"}],
    "consume_rice":        [{"tool": "consume_inventory_item"}],
    "shopping_vegetables": [{"tool": "add_to_shopping_list"}],
    "compare_eggs":        [{"tool": "get_price_history"}],
    "price_tomato":        [{"tool": "record_price_observation"}],
    "use_soon_check":      [{"tool": "check_use_soon"}],
    "buy_suggestions":     [{"tool": "buy_suggestions"}],
    "move_sugar":          [{"tool": "move_inventory_item"}],
    "multi_step":          [{"tool": "add_inventory_item"}],
    "hinglish_add":        [{"tool": "add_inventory_item"}],
    "hinglish_consume":    [{"tool": "consume_inventory_item"}],
    "ambiguous_qty":       [{"tool": "confirm"}, {"tool": "add_inventory_item"}],
    "correction":          [{"tool": "add_inventory_item"}],
    "empty_inventory":     [{"tool": "search_inventory"}],
    "find_multiple":       [{"tool": "search_inventory"}],
    "add_bulk":            [{"tool": "add_inventory_item"}],
    "price_history_check": [{"tool": "get_price_history"}],
    "use_soon_pantry":     [{"tool": "check_use_soon"}],
    "confirm_action":      [{"tool": "confirm"}],
}

# Tool definitions (matching production prompts.py)
TOOL_DESCRIPTIONS = """[
  {{"name": "add_inventory_item", "desc": "Add an item to inventory.", "args": {{"canonical_name": "str", "display_name": "str", "quantity": "float", "unit": "str", "location": "str", "expiry_hint_days": "int"}}}},
  {{"name": "consume_inventory_item", "desc": "Record consumption of an item.", "args": {{"canonical_name": "str", "quantity": "float", "unit": "str"}}}},
  {{"name": "add_to_shopping_list", "desc": "Add item to shopping list.", "args": {{"item_name": "str", "quantity": "float", "unit": "str"}}}},
  {{"name": "search_inventory", "desc": "Search inventory for an item.", "args": {{"query": "str"}}}},
  {{"name": "record_price_observation", "desc": "Record a price observation.", "args": {{"item_name": "str", "price": "float", "store": "str", "unit": "str"}}}},
  {{"name": "get_price_history", "desc": "Get price history for an item.", "args": {{"item_name": "str"}}}},
  {{"name": "check_use_soon", "desc": "Check which items need to be used soon.", "args": {{}}}},
  {{"name": "buy_suggestions", "desc": "Get suggestions for what to buy next.", "args": {{}}}},
  {{"name": "move_inventory_item", "desc": "Move an item to a different location.", "args": {{"canonical_name": "str", "new_location": "str"}}}},
  {{"name": "confirm", "desc": "Ask the user to clarify an ambiguous request.", "args": {{"question": "str"}}}}
]"""

SYSTEM_PROMPT = f"""You are ShopStack, a household shopping memory assistant.
You help track inventory, shopping lists, prices, and freshness for an Indian household.
You respond ONLY with valid JSON tool calls. No prose, no markdown, no explanations.

# Tools
{TOOL_DESCRIPTIONS}

# Rules
- Always respond with a JSON array of tool calls.
- Each tool call has shape: {{"tool": "<name>", "args": {{...}}}}
- If the user request is ambiguous, use the confirm tool to ask for clarification.
- Never invent inventory items not in the request.

# Current inventory
Inventory (empty - new household):
No items currently in inventory.

Shopping list: empty
Recent purchases: none
"""

INVENTORY_CONTEXT = """Inventory (empty - new household):
No items currently in inventory.

Shopping list: empty
Recent purchases: none
"""


def parse_tool_calls(text: str) -> list[dict]:
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    match = re.search(r"\[\s*\{.*?\}\s*\]", text, re.DOTALL)
    if not match:
        match = re.search(r"\{\s*\"tool\".*?\}", text, re.DOTALL)
        if match:
            try:
                obj = json.loads(match.group(0))
                return [obj] if isinstance(obj, dict) else []
            except json.JSONDecodeError:
                return []
        return []
    try:
        result = json.loads(match.group(0))
        return result if isinstance(result, list) else [result]
    except json.JSONDecodeError:
        return []


def score_accuracy(parsed: list[dict], expected: list[dict]) -> bool:
    if not parsed:
        return False
    expected_names = {e["tool"] for e in expected}
    produced_names = {p.get("tool") for p in parsed if isinstance(p, dict)}
    return bool(expected_names & produced_names)


image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch>=2.4.0",
        "transformers>=4.45.0",
        "peft>=0.11.0",
        "accelerate>=0.34.0",
        "bitsandbytes>=0.43.0",
        "huggingface_hub>=0.25.0",
    )
)


@app.function(
    image=image,
    gpu=GPU_CONFIG,
    timeout=600,
    secrets=[modal.Secret.from_name("hf-token")],
    memory=16384,
)
def bench_lora(use_adapter: bool) -> dict:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    label = "LORA" if use_adapter else "BASE-Qwen2.5-1.5B"
    print(f"\n{'='*70}\n[{label}] Loading {'with adapter' if use_adapter else 'base only'}...\n{'='*70}")

    t0 = time.perf_counter()

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=os.environ.get("HF_TOKEN"))
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        token=os.environ.get("HF_TOKEN"),
    )
    if use_adapter:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, LORA_REPO, token=os.environ.get("HF_TOKEN"))
        model = model.merge_and_unload()  # merge for faster inference
    model.eval()

    load_time = time.perf_counter() - t0
    print(f"[{label}] Loaded in {load_time:.1f}s")

    correct = 0
    latencies = []
    details = []

    for pname, question in TEST_PROMPTS:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": question},
        ]
        try:
            input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        except Exception:
            input_text = SYSTEM_PROMPT + "\n\nUser: " + question + "\n\nAssistant:"
        inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

        t0 = time.perf_counter()
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=False,
                temperature=1.0,
                top_p=1.0,
                pad_token_id=tokenizer.eos_token_id,
            )
        latency = time.perf_counter() - t0
        latencies.append(latency)

        generated_ids = outputs[0][inputs.input_ids.shape[1]:]
        output_text = tokenizer.decode(generated_ids, skip_special_tokens=True)

        parsed = parse_tool_calls(output_text)
        is_correct = score_accuracy(parsed, EXPECTED_TOOLS[pname])
        if is_correct:
            correct += 1
        details.append({
            "prompt": pname,
            "correct": is_correct,
            "parsed": [p.get("tool") for p in parsed if isinstance(p, dict)],
            "expected": [e["tool"] for e in EXPECTED_TOOLS[pname]],
            "latency_s": round(latency, 3),
            "output_preview": output_text[:200],
        })

    accuracy = round(100.0 * correct / len(TEST_PROMPTS), 1)
    result = {
        "label": label,
        "accuracy_pct": accuracy,
        "latency_mean_s": round(sum(latencies) / len(latencies), 3),
        "latency_p50_s": round(sorted(latencies)[len(latencies) // 2], 3),
        "load_time_s": round(load_time, 2),
        "details": details,
    }
    print(f"\n[{label}] ✓ {accuracy}% accuracy, {result['latency_mean_s']}s mean")
    return result


@app.local_entrypoint()
def main():
    print("=" * 70)
    print("LoRA ADAPTER TEST — pranaysuyash/shopstack-parser-lora-qwen2.5-1.5b")
    print("=" * 70)
    print(f"Base: {BASE_MODEL}")
    print(f"Adapter: {LORA_REPO}")
    print(f"Test: 20 production PlannerEngine prompts")
    print()

    # Test base + adapter in parallel
    base_result = bench_lora.remote(use_adapter=False)
    lora_result = bench_lora.remote(use_adapter=True)

    # Compare
    print(f"\n{'='*70}")
    print(f"RESULTS COMPARISON")
    print(f"{'='*70}")
    print(f"{'Model':<25} {'Acc':>6} {'Mean Lat':>10} {'Load':>8}")
    print(f"{'-'*70}")
    for r in [base_result, lora_result]:
        print(f"{r['label']:<25} {r['accuracy_pct']:>5}% "
              f"{r['latency_mean_s']:>9.2f}s {r['load_time_s']:>7.1f}s")

    delta = lora_result['accuracy_pct'] - base_result['accuracy_pct']
    print(f"\nLoRA delta: {delta:+.1f}pp")
    if delta > 0:
        print(f"  ✓ LoRA adapter IMPROVES accuracy by {delta:.1f}pp")
    elif delta < 0:
        print(f"  ✗ LoRA adapter HURTS accuracy by {abs(delta):.1f}pp")
    else:
        print(f"  = LoRA adapter is neutral")

    # Save
    out_dir = Path("/Users/pranay/Projects/shopstack/benchmarks/modal/results")
    out_dir.mkdir(parents=True, exist_ok=True)
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    out_file = out_dir / f"lora_test_{timestamp}.jsonl"
    with open(out_file, "w") as f:
        f.write(json.dumps(base_result) + "\n")
        f.write(json.dumps(lora_result) + "\n")
    print(f"\nWrote: {out_file}")

    # Compare to other models
    print(f"\n{'='*70}")
    print(f"COMPARED TO OTHER PLANNER CANDIDATES (from Run 1+2+3)")
    print(f"{'-'*70}")
    print(f"{'Model':<35} {'Acc':>6} {'Params':>8}")
    print(f"  Ministral-8B-Instruct-2410     95.0%  8B   (Run 1+2)")
    print(f"  Gemma-3-4B-it                  95.0%  4B   (Run 2)")
    print(f"  Qwen2.5-7B-Instruct            80.0%  7B   (Run 1+2)")
    print(f"  Qwen3.5-4B (default)           70.0%  4B   (Run 1+2)")
    print(f"  Qwen3.5-4B baseline            70.0%  4B   (Run 3 v3)")
    print(f"  Ministral-3-8B-Reasoning-2512  90.0%  8B   (Run 3)")
    print(f"  Ministral-3-14B-Instruct-2512  85.0%  14B  (Run 3)")
    print(f"  → {base_result['label']:<25} {base_result['accuracy_pct']:>5.0f}%  1.5B")
    print(f"  → {lora_result['label']:<25} {lora_result['accuracy_pct']:>5.0f}%  1.5B + LoRA")
    print(f"{'='*70}")
    if lora_result['accuracy_pct'] >= 90:
        print(f"  ✓✓✓ LoRA adapter is competitive with 8B+ models at 1/5 the size")
        print(f"  This validates the well_tuned badge path.")
    elif lora_result['accuracy_pct'] >= 80:
        print(f"  ✓✓ LoRA adapter is solid (80%+), needs more training data")
    else:
        print(f"  ✓ LoRA adapter trained but needs more data / epochs")