fix bug (#1320)

modelscope · Jul 8, 2024 · 7562138 · 7562138
1 parent f5b4585
commit 7562138
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/swift/llm/utils/preprocess.py b/swift/llm/utils/preprocess.py
@@ -161,7 +161,10 @@ def preprocess(self, d: Dict[str, Any]) -> Dict[str, Any]:
         medias = self.parse_medias(d)
         self.media_replacer(row, medias)
         if self.media_type:
-            row[self.media_name] = medias
+            if not isinstance(self.media_key, str):
+                row[self.media_name] = medias
+            else:
+                row[self.media_key] = medias
         return row
 
     def __call__(self, dataset: HfDataset) -> HfDataset:
@@ -251,7 +254,10 @@ def preprocess(self, d: Dict[str, Any]) -> Dict[str, Any]:
             medias = self.parse_medias(d)
             self.media_replacer(row, medias)
             if self.media_type:
-                row[self.media_name] = medias
+                if not isinstance(self.media_key, str):
+                    row[self.media_name] = medias
+                else:
+                    row[self.media_key] = medias
             return row
         except (AssertionError, SyntaxError):
             if self.error_strategy == 'raise':
@@ -306,7 +312,10 @@ def preprocess(self, d: Dict[str, Any]) -> Dict[str, Any]:
             medias = self.parse_medias(d)
             self.media_replacer(row, medias)
             if self.media_type:
-                row[self.media_name] = medias
+                if not isinstance(self.media_key, str):
+                    row[self.media_name] = medias
+                else:
+                    row[self.media_key] = medias
         except Exception:
             if self.error_strategy == 'raise':
                 raise ValueError(f'conversations: {conversations}')