diff --git a/modeling_deepseekocr.py b/modeling_deepseekocr.py index 05ebf94..b39f210 100644 --- a/modeling_deepseekocr.py +++ b/modeling_deepseekocr.py @@ -944,17 +944,6 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM): no_repeat_ngram_size = 35, use_cache = True ) - - - if '' 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 '' in conversation[0]['content'] and test_compress: outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:]) @@ -1035,3 +1024,13 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM): plt.close() result.save(f"{output_path}/result_with_boxes.jpg") + + if '' 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