electblake commited on
Commit
25c4ebd
·
1 Parent(s): 8bb200e

fix(model): parse JSON after thinking output

Browse files

Handle 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.

Files changed (1) hide show
  1. model.py +7 -5
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
- processor.batch_decode(
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