DESC: fix infer

This commit is contained in:
xiawei 2025-11-14 15:43:41 +08:00
parent 252e9b59e1
commit cb95a394a5

View File

@ -945,17 +945,6 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
use_cache = True use_cache = True
) )
if '<image>' in conversation[0]['content'] and eval_mode:
outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
stop_str = '<end▁of▁sentence>'
if outputs.endswith(stop_str):
outputs = outputs[:-len(stop_str)]
# re_match
outputs = outputs.strip()
return outputs
if '<image>' in conversation[0]['content'] and test_compress: if '<image>' in conversation[0]['content'] and test_compress:
outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:]) outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
pure_texts_outputs_token_length = len(text_encode(tokenizer, outputs, bos=False, eos=False)) pure_texts_outputs_token_length = len(text_encode(tokenizer, outputs, bos=False, eos=False))
@ -1035,3 +1024,13 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
plt.close() plt.close()
result.save(f"{output_path}/result_with_boxes.jpg") result.save(f"{output_path}/result_with_boxes.jpg")
if '<image>' in conversation[0]['content'] and eval_mode:
outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
stop_str = '<end▁of▁sentence>'
if outputs.endswith(stop_str):
outputs = outputs[:-len(stop_str)]
# re_match
outputs = outputs.strip()
return outputs