Spaces:
Running on Zero
Running on Zero
electblake commited on
Commit ·
25c4ebd
1
Parent(s): 8bb200e
fix(model): parse JSON after thinking output
Browse filesHandle NuExtract3's documented thinking response format by decoding the final answer after the closing think tag before JSON parsing. Non-thinking extraction remains unchanged and malformed thinking output still fails at its source.
model.py
CHANGED
|
@@ -53,12 +53,14 @@ def extract(image, text, template, enable_thinking):
|
|
| 53 |
do_sample=False,
|
| 54 |
)
|
| 55 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
return json.loads(
|
| 57 |
-
|
| 58 |
-
generated_ids[:, inputs.input_ids.shape[1] :],
|
| 59 |
-
skip_special_tokens=True,
|
| 60 |
-
clean_up_tokenization_spaces=False,
|
| 61 |
-
)[0].strip()
|
| 62 |
)
|
| 63 |
|
| 64 |
|
|
|
|
| 53 |
do_sample=False,
|
| 54 |
)
|
| 55 |
|
| 56 |
+
output = processor.batch_decode(
|
| 57 |
+
generated_ids[:, inputs.input_ids.shape[1] :],
|
| 58 |
+
skip_special_tokens=True,
|
| 59 |
+
clean_up_tokenization_spaces=False,
|
| 60 |
+
)[0].strip()
|
| 61 |
+
|
| 62 |
return json.loads(
|
| 63 |
+
output.split("</think>", 1)[1].strip() if enable_thinking else output
|
|
|
|
|
|
|
|
|
|
|
|
|
| 64 |
)
|
| 65 |
|
| 66 |
|