VaidikML0508/qwen2.5-3b-investigation-grpo

pip install -U transformers peft accelerate bitsandbytes torch
# optional (private repos / gated base):
pip install -U huggingface_hub
huggingface-cli login

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

BASE_MODEL = "Qwen/Qwen2.5-3B-Instruct"
ADAPTER_ID = "VaidikML0508/mystery-rl-adapter22"  # or mystery-sft-adapter22

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=bnb,
    device_map="auto",
    trust_remote_code=True,
)
model = PeftModel.from_pretrained(base, ADAPTER_ID)
model.eval()

base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)
model = PeftModel.from_pretrained(base, ADAPTER_ID)

<tool_call>{"name":"call_action","arguments":{"action_name":"inspect_study","arguments_json":"{}"}}</tool_call>

<tool_call>{"name":"close_case","arguments":{"culprit":"Jordan Lee","citations_json":"[\"ev_ledger\",\"ev_key\"]"}}</tool_call>

import re
import torch

SYSTEM = (
    "You are an investigation agent. On EVERY turn output exactly ONE tool call, then stop.\n"
    "Format (exact):\n"
    '<tool_call>{"name":"call_action","arguments":{"action_name":"ACTION","arguments_json":"{}"}}</tool_call>\n'
    "or\n"
    '<tool_call>{"name":"close_case","arguments":{"culprit":"NAME","citations_json":"[\\"ev1\\",\\"ev2\\"]"}}</tool_call>\n'
    "Rules:\n"
    "- Only use currently unlocked action names from the user message.\n"
    "- Do not repeat the same action if it already returned evidence.\n"
    "- After you have required evidence, call close_case with citations you observed.\n"
)

USER = """Case brief: A rare manuscript vanished from the university archive overnight.
Suspects: Avery Quinn, Jordan Lee, Sam Ortiz
Places: archive, loading dock, dormitory

Unlocked actions: list_archive_log, interview_night_guard
Discovered evidence: []

Call exactly one tool."""

messages = [
    {"role": "system", "content": SYSTEM},
    {"role": "user", "content": USER},
]

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=160,
        do_sample=True,
        temperature=0.3,
        top_p=0.9,
        pad_token_id=tokenizer.pad_token_id,
    )

gen = tokenizer.decode(out[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)
print(gen)

# Optional: extract first tool_call JSON
m = re.search(r"<tool_call>\s*(\{.*?\})\s*(?:</tool_call>|$)", gen, flags=re.DOTALL)
print("parsed:", m.group(1) if m else None)

import json
import re
from typing import Any

def extract_tool_call(text: str) -> dict[str, Any] | None:
    m = re.search(
        r"<tool_call>\s*(\{.*?\})\s*(?:</tool_call>|(?=<tool_call>)|$)",
        text,
        flags=re.DOTALL,
    )
    if not m:
        return None
    try:
        return json.loads(m.group(1))
    except json.JSONDecodeError:
        return None


def step_env(tool: dict[str, Any]) -> str:
    """Replace with your InvestigationGym / API."""
    name = tool.get("name")
    args = tool.get("arguments") or {}
    if name == "call_action":
        return env.call_action(args["action_name"], args.get("arguments_json", "{}"))
    if name == "close_case":
        return env.close_case(args["culprit"], args["citations_json"])
    return f"ERROR: unknown tool {name}"


messages = [
    {"role": "system", "content": SYSTEM},
    {"role": "user", "content": initial_case_prompt},  # brief + unlocked + discovered
]

for turn in range(12):
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=160,
            do_sample=True,
            temperature=0.3,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id,
        )
    gen = tokenizer.decode(out[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)

    # Keep only the first tool_call block if the model continues
    if "<tool_call>" in gen:
        i = gen.find("<tool_call>")
        rest = gen[i:]
        j = rest.find("</tool_call>")
        gen = rest[: j + len("</tool_call>")] if j != -1 else rest

    messages.append({"role": "assistant", "content": gen})
    tool = extract_tool_call(gen)
    if not tool:
        messages.append({
            "role": "user",
            "content": "ERROR: no valid tool_call. Output exactly one <tool_call>{...}</tool_call>.",
        })
        continue

    obs = step_env(tool)
    # Include updated unlocked / discovered lists in the user turn when you have them
    messages.append({
        "role": "user",
        "content": (
            f"Tool result ({tool.get('name')}):\n{obs}\n"
            "Continue with exactly one next tool_call "
            "(new evidence action, or close_case if ready)."
        ),
    })

    if tool.get("name") == "close_case" and not str(obs).startswith("ERROR"):
        break