Allow uploading tool calling datasets

Author: connermanuel

src/together/lib/cli/api/_utils.py

@@ -160,6 +160,10 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
             except Exception as e:
                 click.echo(prefix_styled + click.style("Failed", fg="red"))
                 click.echo(prefix_styled + click.style(f"An unexpected error occurred - {str(e)}", fg="red"))
+                import traceback
+
+                click.echo(prefix_styled + click.style("Traceback:", fg="red"))
+                click.echo(prefix_styled + click.style(traceback.format_exc(), fg="red"))
                 sys.exit(1)
 
         return wrapper  # type: ignore

src/together/lib/constants.py

@@ -66,5 +66,11 @@ class DatasetFormat(enum.Enum):
         "non_preferred_output",
     ],
 }
-REQUIRED_COLUMNS_MESSAGE = ["role", "content"]
+JSONL_ADDITIONAL_COLUMNS_MAP = {
+    DatasetFormat.GENERAL: [],
+    DatasetFormat.CONVERSATION: ["tools"],
+    DatasetFormat.INSTRUCTION: [],
+    DatasetFormat.PREFERENCE_OPENAI: [],
+}
+REQUIRED_COLUMNS_MESSAGE = ["role"]
 POSSIBLE_ROLES_CONVERSATION = ["system", "user", "assistant"]

src/together/lib/utils/files.py

@@ -22,6 +22,7 @@
     REQUIRED_COLUMNS_MESSAGE,
     JSONL_REQUIRED_COLUMNS_MAP,
     POSSIBLE_ROLES_CONVERSATION,
+    JSONL_ADDITIONAL_COLUMNS_MAP,
     DatasetFormat,
 )
 
@@ -392,7 +393,14 @@ def validate_messages(
         message_weight = _check_message_weight(message, idx)
         previous_role = _check_message_role(message, previous_role, idx)
         assistant_role_exists |= previous_role == "assistant"
-        is_multimodal, number_of_images = _check_message_content(message["content"], role=previous_role, idx=idx)
+
+        number_of_images = 0
+        is_multimodal = False
+
+        content = message.get("content")
+        if content is not None:
+            is_multimodal, number_of_images = _check_message_content(content, role=previous_role, idx=idx)
+
         # Multimodal validation
         if number_of_images > 0 and message_weight is not None and message_weight != 0:
             raise InvalidFileFormatError(
@@ -656,7 +664,11 @@ def _check_jsonl(file: Path, purpose: FilePurpose | str) -> Dict[str, Any]:
 
                             # Check that there are no extra columns
                             for column in cast(List[str], json_line.keys()):
-                                if column not in JSONL_REQUIRED_COLUMNS_MAP[possible_format]:
+                                allowed_columns = (
+                                    JSONL_REQUIRED_COLUMNS_MAP[possible_format]
+                                    + JSONL_ADDITIONAL_COLUMNS_MAP[possible_format]
+                                )
+                                if column not in allowed_columns:
                                     raise InvalidFileFormatError(
                                         message=f'Found extra column "{column}" in the line {idx + 1}.',
                                         line_number=idx + 1,