Views
No views yet
Author · Seungyoun Shin🤗 Model Hub: hf
nq-hotpotqa-train with GRPO via the open‑source VERL framework.pip install "transformers>=4.41" torch duckduckgo_search>=6.3.5 accelerate1#!/usr/bin/env python3
2"""
3Minimal **multi‑turn tool‑calling** demo for the Qwen2.5‑3b‑it_searchR1‑like model
4
5Key points
6-----------
7* Supplies the DuckDuckGo *search* tool schema via `tools=[…]` so the model emits JSON‑style calls.
8* Detects `<tool_call>` → parses JSON `{name:…, arguments:{query_list:[…]}}` and runs DuckDuckGo for each query.
9* Streams the results back inside `<tool_response>` so the model can reason again, up to `MAX_TURNS`.
10
11Install once:
12 pip install "duckduckgo_search>=6.3.5"
13
14Run:
15 python3 search_r1_infer.py "How is the weather in Seoul?"
16"""
17
18from __future__ import annotations
19
20import json
21import re
22import sys
23from typing import List
24
25import torch
26from transformers import AutoModelForCausalLM, AutoTokenizer
27from duckduckgo_search import DDGS
28
29# ----------------------------------------------------------------------------
30# Color codes for terminal output
31# ----------------------------------------------------------------------------
32class Colors:
33 RESET = '\033[0m'
34 BOLD = '\033[1m'
35 RED = '\033[91m'
36 GREEN = '\033[92m'
37 YELLOW = '\033[93m'
38 BLUE = '\033[94m'
39 MAGENTA = '\033[95m'
40 CYAN = '\033[96m'
41
42# ----------------------------------------------------------------------------
43# Constants & Prompt Template
44# ----------------------------------------------------------------------------
45DEFAULT_SYSTEM_CONTENT = "You are a helpful and harmless assistant."
46DEFAULT_USER_CONTENT_PREFIX = (
47 "Answer the given question. You must conduct reasoning inside <think> and "
48 "</think> first every time you get new information. After reasoning, if you "
49 "find you lack some knowledge, you can call a search engine by <tool_call> "
50 "query </tool_call> and it will return the top searched results between "
51 "<tool_response> and </tool_response>. You can search as many times as your "
52 "want. If you find no further external knowledge needed, you can directly "
53 "provide the answer inside <answer> and </answer>, without detailed "
54 "illustrations. For example, <answer> Beijing </answer>. Question: "
55)
56
57MODEL_NAME = "Seungyoun/qwen2.5-7b-it_searchR1-like-sgl-multiturn"
58MAX_TURNS = 5
59MAX_RESPONSE_TOKENS = 8192
60DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
61
62# ----------------------------------------------------------------------------
63# Tool schema (JSON mirror of search_tool_config.yaml)
64# ----------------------------------------------------------------------------
65SEARCH_SCHEMA = {
66 "type": "function",
67 "function": {
68 "name": "search",
69 "description": "Searches the web for relevant information based on the given query.",
70 "parameters": {
71 "type": "object",
72 "properties": {
73 "query_list": {
74 "type": "array",
75 "items": {"type": "string"},
76 "description": (
77 "A list of fully‑formed semantic queries. The tool will return "
78 "search results for each query."
79 ),
80 }
81 },
82 "required": ["query_list"],
83 },
84 },
85}
86
87# ----------------------------------------------------------------------------
88# Helper functions
89# ----------------------------------------------------------------------------
90
91def create_prompt(question: str) -> List[dict]:
92 """Build the initial chat prompt."""
93 return [
94 {"role": "system", "content": DEFAULT_SYSTEM_CONTENT},
95 {"role": "user", "content": DEFAULT_USER_CONTENT_PREFIX + question},
96 ]
97
98
99def ddg_search_one(query: str, k: int = 5) -> str:
100 """Return top‑k DuckDuckGo results joined by newlines."""
101 with DDGS() as ddgs:
102 hits = list(ddgs.text(query, safesearch="moderate", max_results=k))
103 return "\n".join(
104 f"{i+1}. {h['title']} – {h['body']} ({h['href']})" for i, h in enumerate(hits)
105 )
106
107
108def extract_queries(raw: str) -> List[str]:
109 """Parse the JSON inside <tool_call> and return the `query_list`. Fallback to raw."""
110 try:
111 payload = json.loads(raw)
112 if (
113 isinstance(payload, dict)
114 and payload.get("name") == "search"
115 and isinstance(payload.get("arguments"), dict)
116 ):
117 qlist = payload["arguments"].get("query_list", [])
118 return [q for q in qlist if isinstance(q, str)]
119 except json.JSONDecodeError:
120 pass # raw is not JSON → treat as literal
121 return [raw]
122
123
124# ----------------------------------------------------------------------------
125# Main driver
126# ----------------------------------------------------------------------------
127
128def main() -> None:
129 question = sys.argv[1] if len(sys.argv) > 1 else "Who is the president of Korea?"
130
131 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side="left")
132 model = AutoModelForCausalLM.from_pretrained(
133 MODEL_NAME, torch_dtype=torch.bfloat16, device_map="auto"
134 )
135
136 messages = create_prompt(question)
137 chat_history = tokenizer.apply_chat_template(
138 messages,
139 tools=[SEARCH_SCHEMA], # expose tool to the model
140 add_generation_prompt=True,
141 tokenize=False,
142 )
143
144 tool_call_pattern = re.compile(r"<tool_call>\s*(.*?)\s*</tool_call>", re.S)
145
146 for turn in range(MAX_TURNS):
147 chat_history = tokenizer.apply_chat_template(
148 messages,
149 tools=[SEARCH_SCHEMA], # expose tool to the model
150 add_generation_prompt=True,
151 tokenize=False,
152 )
153 enc = tokenizer(chat_history, return_tensors="pt").to(DEVICE)
154 out = model.generate(
155 **enc,
156 max_new_tokens=MAX_RESPONSE_TOKENS,
157 temperature=0.1,
158 top_p=0.9,
159 do_sample=True,
160 )
161 new_text = tokenizer.decode(out[0][enc.input_ids.shape[1] :], skip_special_tokens=True)
162 print(f"\n===== Assistant (turn {turn+1}) =====\n{new_text}\n")
163 chat_history += new_text
164
165 m = tool_call_pattern.search(new_text)
166 if not m:
167 break # finished – no tool call
168
169 queries = extract_queries(m.group(1))
170 all_results: list[str] = []
171 for q in queries:
172 print(f"{Colors.CYAN}{Colors.BOLD}[Tool Call] 검색 쿼리: {q}{Colors.RESET}")
173 search_result = ddg_search_one(q, k=5)
174 all_results.append(search_result)
175 print(f"{Colors.GREEN}[Tool Response]{Colors.RESET}")
176 print(f"{Colors.GREEN}{search_result}{Colors.RESET}")
177 print(f"{Colors.GREEN}{'='*50}{Colors.RESET}\n")
178
179 tool_response_block = "<tool_response>\n" + "\n---\n".join(all_results) + "\n</tool_response>"
180 messages.append({"role": "user", "content": tool_response_block})
181 chat_history += tool_response_block # feed back into next turn
182
183if __name__ == "__main__":
184 main()<think> … chain‑of‑thought … </think>
<tool_call>{"name":"search", "arguments":{"query_list":["…"]}}</tool_call>
<tool_response>
1. web result
…
</tool_response>
<answer> final concise answer </answer>| Dataset | Original Search‑R1 (Qwen2.5‑7B-it) | This work |
|---|---|---|
| NQ | 0.393 | 0.422 |
| TriviaQA | 0.610 | 0.612 |
| PopQA | 0.397 | 0.450 |
| HotpotQA | 0.370 | 0.366 |
| 2Wiki | 0.414 | 0.371 |
| Musique | 0.143 | 0.124 |
| Bamboogle | 0.368 | 0.376 |
| Avg. | 0.385 | 0.389 |
1@misc{shin2025qwen25_searchr1_multiturn,
2 author = {Seungyoun Shin},
3 title = {Qwen2.5-3B Search-R1-Multiturn (reproduce)},
4 year = 2025,
5 howpublished = {HuggingFace Model Hub},
6 url = {https://huggingface.co/Seungyoun/qwen2.5-3b-it_searchR1-like-multiturn}
7}