1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model = AutoModelForCausalLM.from_pretrained("caraman/Qwen2.5-7B-mtrag-query-rewriter-final")
4tokenizer = AutoTokenizer.from_pretrained("caraman/Qwen2.5-7B-mtrag-query-rewriter-final")
5
6system_prompt = """You are a query rewriting assistant for information retrieval. Given a conversation history and a current question, rewrite the question to be completely standalone and self-contained.
7
8Rules:
91. Resolve all pronouns (it, they, this, that) to their explicit referents
102. Include relevant context from the conversation that's needed to understand the query
113. Keep the rewritten query concise and search-friendly
124. Do not add information not present in the conversation
135. If the question is already standalone, return it unchanged"""
14
15messages = [
16 {"role": "system", "content": system_prompt},
17 {"role": "user", "content": """CONVERSATION HISTORY:
18USER: Tell me about the Eiffel Tower
19ASSISTANT: The Eiffel Tower is a wrought-iron lattice tower in Paris, France.
20
21CURRENT QUESTION: When was it built?
22
23Rewrite this question to be standalone:"""}
24]
25
26text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
27inputs = tokenizer(text, return_tensors="pt")
28outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.2)
29print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
30# Expected: "When was the Eiffel Tower built?"
1from mlx_lm import load, generate
2
3model, tokenizer = load("caraman/Qwen2.5-7B-mtrag-query-rewriter-final")
4prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
5response = generate(model, tokenizer, prompt=prompt, max_tokens=256, temp=0.2)
Part of a three-stage pipeline (query rewriting + hybrid BM25/dense retrieval + cross-encoder reranking):