import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
BASE = "Qwen/Qwen2.5-0.5B-Instruct"
ADAPTER = "eorgantzoglou/qwen2.5-0.5b-airline-triage-lora"
SYSTEM_PROMPT = (
"You are a triage assistant for an airline's customer support. "
"Classify the customer tweet. Respond with json only, in exactly this format: "
'{"intent": "...", "urgency": "...", "abusive": true/false}. '
"intent must be one of: delay_disruption, checkin_boarding_issue, "
"flight_cancellation_rebooking, lost_luggage, special_assistance, "
"general_complaint, general_question, praise_feedback, spam_irrelevant, "
"other_unclear. urgency must be one of: high, medium, low."
)
tokenizer = AutoTokenizer.from_pretrained(BASE)
model = PeftModel.from_pretrained(
AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.float32), ADAPTER
).merge_and_unload()
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "my bags are lost and nobody at the desk is helping"},
]
inputs = tokenizer(
tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True),
return_tensors="pt",
)
out = model.generate(**inputs, max_new_tokens=60, do_sample=False,
pad_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))