fix unirec postprocess

Author: topdu

demo_unirec.py

@@ -3,7 +3,6 @@
 from threading import Thread
 
 import numpy as np
-import re
 from openrec.postprocess.unirec_postprocess import clean_special_tokens
 from openrec.preprocess import create_operators, transform
 from tools.engine.config import Config
@@ -41,18 +40,6 @@ def set_device(device):
 transforms, ratio_resize_flag = build_rec_process(cfg)
 ops = create_operators(transforms, global_config)
 
-rules = [
-    (r'-<\|sn\|>', ''),
-    (r' <\|sn\|>', ' '),
-    (r'<\|sn\|>', ' '),
-    (r'<\|unk\|>', ''),
-    (r'<s>', ''),
-    (r'</s>', ''),
-    (r'\uffff', ''),
-    (r'_{4,}', '___'),
-    (r'\.{4,}', '...'),
-]
-
 
 # --- 2. 定义流式生成函数 ---
 def stream_chat_with_image(input_image, history):
@@ -74,17 +61,20 @@ def stream_chat_with_image(input_image, history):
         'input_ids': None,
         'attention_mask': None
     }
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048)
     # 后台线程运行生成
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     # 流式输出
-    generated_text = ''
     history = history + [('🖼️(图片)', '')]
+    generated_text_ori = ''
     for new_text in streamer:
-        generated_text += clean_special_tokens(new_text)
-        for rule in rules:
-            generated_text = re.sub(rule[0], rule[1], generated_text)
+        generated_text_ori += new_text
+        generated_text = clean_special_tokens(
+            generated_text_ori.replace(' ', ''))
+        text = generated_text.replace('<tdcolspan=', '<td colspan=')
+        text = text.replace('<tdrowspan=', '<td rowspan=')
+        generated_text = text.replace('"colspan=', '" colspan=')
         history[-1] = ('🖼️(图片)', generated_text)
         yield history
 

openrec/postprocess/unirec_postprocess.py

@@ -15,13 +15,11 @@
 
 
 def clean_special_tokens(text):
-    text = text.replace(' ', '').replace('Ġ', ' ').replace('Ċ', '\n').replace(
-        '<|bos|>', '').replace('<|eos|>', '').replace('<|pad|>', '')
+    text = text.replace('Ġ',
+                        ' ').replace('Ċ', '\n').replace('<|bos|>', '').replace(
+                            '<|eos|>', '').replace('<|pad|>', '')
     for rule in rules:
         text = re.sub(rule[0], rule[1], text)
-    text = text.replace('<tdcolspan=', '<td colspan=')
-    text = text.replace('<tdrowspan=', '<td rowspan=')
-    text = text.replace('"colspan=', '" colspan=')
     return text
 
 
@@ -44,7 +42,10 @@ def __init__(self,
     def __call__(self, preds, batch=None, *args, **kwargs):
         result_list = []
         pred_ids = preds
-        res = self.tokenizer.batch_decode(pred_ids, skip_special_tokens=False)
+        res = [
+            ''.join(self.tokenizer.convert_ids_to_tokens(seq.tolist()))
+            for seq in pred_ids
+        ]
         for i in range(len(res)):
             res[i] = clean_special_tokens(res[i])
             result_list.append(