fix modeling_sa2va_chat.py bug

modeling_sa2va_chat.py

@@ -807,7 +807,7 @@ class Sa2VAChatModel(PreTrainedModel):
         bg_cls_token_id = torch.as_tensor([self.bg_cls_token_idx,], dtype=ids.dtype, device=ids.device)
         bg_cls_embedding = self.language_model.get_input_embeddings()(bg_cls_token_id).clone()
         output_ids = generate_output.sequences[0][:-1]
-        cls_token_mask = ids == self.cls_token_idx
+        cls_token_mask = ids[0] == self.cls_token_idx
 
         # get seg tokens
         seg_token_mask = (output_ids >= self.the_first_seg_token_idx) & (output_ids <= self.the_last_seg_token_idx)
@@ -841,7 +841,7 @@ class Sa2VAChatModel(PreTrainedModel):
                 **m2f_inputs
             )
 
-            tags = re.findall(r'<p>(.*?)</p>', predict)
+            tags = re.findall(r'<p>(.*?)</p>', input_text)
             label_id_to_text = {id: tag for id, tag in enumerate(tags)}
             
             class_queries_logits = m2f_outputs.class_queries_logits