Merge pull request #38 from Cmxu/collators_fix

fix system prompts
zjysteven · Oct 18, 2024 · b4ee72c · b4ee72c
2 parents 5db747c + 816c126
commit b4ee72c
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/collators/llava_1_5.py b/collators/llava_1_5.py
@@ -51,7 +51,7 @@ def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
             if system_prompt is not None:
                 cur_text.append({
                     "role": "system",
-                    "content": [{"text": system_prompt}]
+                    "content": [{"type": "text", "text": system_prompt}]
                 })
 
             for i, text in enumerate(cur_convs):

diff --git a/collators/llava_1_6.py b/collators/llava_1_6.py
@@ -51,7 +51,7 @@ def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
             if system_prompt is not None:
                 cur_text.append({
                     "role": "system",
-                    "content": [{"text": system_prompt}]
+                    "content": [{"type": "text", "text": system_prompt}]
                 })
 
             for i, text in enumerate(cur_convs):

diff --git a/collators/llava_next_video.py b/collators/llava_next_video.py
@@ -59,7 +59,7 @@ def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
             if system_prompt is not None:
                 cur_text.append({
                     "role": "system",
-                    "content": [{"text": system_prompt}]
+                    "content": [{"type": "text", "text": system_prompt}]
                 })
 
             for i, text in enumerate(cur_convs):

diff --git a/collators/qwen2_vl.py b/collators/qwen2_vl.py
@@ -61,7 +61,7 @@ def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
             if system_prompt is not None:
                 cur_text.append({
                     "role": "system",
-                    "content": [{"text": system_prompt}]
+                    "content": [{"type": "text", "text": system_prompt}]
                 })
 
             for i, text in enumerate(cur_convs):