Views
No views yet
| Property | Value |
|---|---|
| Base model | Qwen/Qwen2.5-1.5B-Instruct |
| Fine-tuning method | QLoRA (4-bit NF4 + LoRA) |
| LoRA rank | 16 |
| LoRA alpha | 32 |
| Target modules | q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj |
| Training examples | 481 (90% of 535 total) |
| Eval examples | 54 (10% of 535 total) |
| Training epochs | 3 |
| Effective batch size | 16 (4 × 4 gradient accumulation) |
| Learning rate | 2e-4 (cosine schedule) |
| Max sequence length | 256 |
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4model_id = "abi-commits/qwen-query-optimizer"
5tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
6model = AutoModelForCausalLM.from_pretrained(
7 model_id,
8 device_map="auto",
9 torch_dtype=torch.float16,
10 trust_remote_code=True,
11)
12
13SYSTEM_PROMPT = (
14 "You are a query optimization agent. Rewrite user queries into clear, "
15 "retrieval-focused enterprise document search queries. "
16 "Do not add new information. Do not hallucinate."
17)
18
19def optimize_query(user_query: str) -> str:
20 messages = [
21 {"role": "system", "content": SYSTEM_PROMPT},
22 {"role": "user", "content": user_query},
23 ]
24 prompt = tokenizer.apply_chat_template(
25 messages, tokenize=False, add_generation_prompt=True
26 )
27 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
28 with torch.no_grad():
29 output_ids = model.generate(
30 **inputs,
31 max_new_tokens=80,
32 do_sample=False,
33 repetition_penalty=1.1,
34 eos_token_id=tokenizer.eos_token_id,
35 pad_token_id=tokenizer.pad_token_id,
36 )
37 generated = output_ids[0][inputs["input_ids"].shape[1]:]
38 return tokenizer.decode(generated, skip_special_tokens=True).strip()
39
40# Examples
41print(optimize_query("how do i request time off?"))
42# → "employee leave request procedure and time-off policy"
43
44print(optimize_query("what's the refund policy?"))
45# → "refund policy terms and conditions for customer returns"