Fix Qwen3 thinking mode + increase max_new_tokens: inference.py
Browse files- inference.py +10 -1
inference.py
CHANGED
|
@@ -20,6 +20,7 @@ Environment variables (mandatory):
|
|
| 20 |
import asyncio
|
| 21 |
import json
|
| 22 |
import os
|
|
|
|
| 23 |
import sys
|
| 24 |
import time
|
| 25 |
from typing import List, Optional
|
|
@@ -159,9 +160,17 @@ def build_observation_prompt(obs, is_first: bool = False) -> str:
|
|
| 159 |
# -- LLM response parsing ------------------------------------------------
|
| 160 |
|
| 161 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 162 |
def parse_llm_response(response_text: str) -> dict:
|
| 163 |
"""Extract a JSON object from the LLM response."""
|
| 164 |
-
text = response_text.strip()
|
| 165 |
|
| 166 |
if "```json" in text:
|
| 167 |
text = text.split("```json")[1].split("```")[0].strip()
|
|
|
|
| 20 |
import asyncio
|
| 21 |
import json
|
| 22 |
import os
|
| 23 |
+
import re
|
| 24 |
import sys
|
| 25 |
import time
|
| 26 |
from typing import List, Optional
|
|
|
|
| 160 |
# -- LLM response parsing ------------------------------------------------
|
| 161 |
|
| 162 |
|
| 163 |
+
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def strip_think_blocks(text: str) -> str:
|
| 167 |
+
"""Remove Qwen3-style <think>...</think> reasoning blocks."""
|
| 168 |
+
return _THINK_RE.sub("", text).strip()
|
| 169 |
+
|
| 170 |
+
|
| 171 |
def parse_llm_response(response_text: str) -> dict:
|
| 172 |
"""Extract a JSON object from the LLM response."""
|
| 173 |
+
text = strip_think_blocks(response_text).strip()
|
| 174 |
|
| 175 |
if "```json" in text:
|
| 176 |
text = text.split("```json")[1].split("```")[0].strip()
|