import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "TheArchitect256/qwen3.5-2b-triage-master"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
{}
### Input:
{}
### Response:
"""
instruction = """Available tools:
1. [STOCK_CHECK]: Verifies if an item is available in the warehouse inventory.
2. [SHIPPING_CALC]: Calculates delivery rates and times based on postal codes.
User Query: Can you tell me if we still have the RTX 4060 graphics card in our warehouse?"""
prompt = alpaca_prompt.format(instruction, "")
inputs = tokenizer(text=prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=20, do_sample=False, pad_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True))
import re
match = re.search(r"\[[A-Z_]+\]", response)
clean_response = match.group(0) if match else response.strip()