Update modeling_llada2uni_moe.py

modeling_llada2uni_moe.py

@@ -2419,9 +2419,12 @@ class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
         tok = self._get_tokenizer(tokenizer)
         sp = self._get_special_tokens(tok, image_h, image_w)
 
-        img_header = self._build_image_header(sp)
-        pfx = tok(question).input_ids if question else []
-        ids = img_header + image_tokens + sp["eoi"] + pfx
+        user = self._build_image_header(sp) + image_tokens + sp["eoi"] \
+            + tok("\n").input_ids + (tok(question).input_ids if question else [])
+        sys_ids, user_ids, asst_ids = self._build_chat(
+            tok, "You are a multimodal understanding assistant.", user,
+        )
+        ids = sys_ids + user_ids + asst_ids
 
         out = self.generate_bd(
             data={"input_ids": torch.tensor(ids).unsqueeze(0).to(self.device)},
@@ -2433,7 +2436,7 @@ class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
             image_keep_ratio=image_keep_ratio, text_keep_ratio=text_keep_ratio,
             show_progress=False,
         )
-        return tok.decode(out[0][len(ids) - len(pfx):], skip_special_tokens=True)
+        return tok.decode(out[0][len(ids):], skip_special_tokens=True)
 
     @torch.no_grad()
     def edit_image(self, image_tokens, image_h, image_w, instruction,