# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from pathlib import Path
from typing import Any, Optional, Union
import pandas as pd
from typing_extensions import override
from oumi.core.datasets import BaseSftDataset
from oumi.core.registry import register_dataset
from oumi.core.types.conversation import Conversation
from oumi.utils.io_utils import load_json, load_jsonlines
[docs]
@register_dataset("text_sft")
@register_dataset("text_sft_jsonl")
class TextSftJsonLinesDataset(BaseSftDataset):
"""TextSftJsonLinesDataset for loading SFT data in oumi and alpaca formats.
This dataset class is designed to work with JSON Lines (.jsonl) or
JSON (.json) files containing text-based supervised fine-tuning (SFT) data.
It supports loading data either from a file or from a provided list of data
samples in oumi and alpaca formats.
Supported formats:
1. JSONL or JSON of conversations (Oumi format)
2. JSONL or JSON of Alpaca-style turns (instruction, input, output)
Args:
dataset_path (Optional[Union[str, Path]]): Path to the dataset file
(.jsonl or .json).
data (Optional[List[Dict[str, Any]]]): List of conversation dicts if not
loading from a file.
format (Optional[str]): The format of the data. Either "conversations" or
"alpaca". If not provided, the format will be auto-detected.
**kwargs: Additional arguments to pass to the parent class.
Examples:
Loading conversations from a JSONL file with auto-detection:
>>> from oumi.datasets import TextSftJsonLinesDataset
>>> dataset = TextSftJsonLinesDataset( # doctest: +SKIP
... dataset_path="/path/to/your/dataset.jsonl"
... )
Loading Alpaca-style data from a JSON file:
>>> from oumi.datasets import TextSftJsonLinesDataset
>>> dataset = TextSftJsonLinesDataset( # doctest: +SKIP
... dataset_path="/path/to/your/dataset.json",
... format="alpaca"
... )
Loading from a list of data samples:
>>> from oumi.datasets import TextSftJsonLinesDataset
>>> data_samples = [
... {"messages": [{"role": "user", "content": "Hello"},
... {"role": "assistant", "content": "Hi there!"}]},
... {"messages": [{"role": "user", "content": "How are you?"},
... {"role": "assistant", "content": "great!"}]}
... ]
>>> dataset = TextSftJsonLinesDataset(
... data=data_samples,
... )
"""
default_dataset = "custom"
def __init__(
self,
dataset_path: Optional[Union[str, Path]] = None,
data: Optional[list[dict[str, Any]]] = None,
format: Optional[str] = None,
**kwargs,
):
"""Initializes a new instance of the TextSftJsonLinesDataset class.
Args:
dataset_path (Optional): Path to the JSON lines dataset file.
data (Optional): List of conversation dicts if not loading from a file.
format (Optional): The format of the data. Either "conversations",
or "alpaca". If not provided, the format will be
auto-detected.
**kwargs: Additional arguments to pass to the parent class.
Raises:
ValueError: If neither dataset_path nor data is provided,
or if both are provided.
"""
if dataset_path is not None and data is not None:
raise ValueError(
"Either dataset_path or data must be provided, but not both"
)
self._data_column: str = "_messages_column"
self._dataset_path: Optional[Path] = (
Path(dataset_path) if dataset_path else None
)
if data is not None:
data_frame = pd.DataFrame({self._data_column: data})
elif self._dataset_path is not None:
if self._dataset_path.suffix.lower() == ".jsonl":
data = load_jsonlines(self._dataset_path)
elif self._dataset_path.suffix.lower() == ".json":
data = load_json(self._dataset_path)
else:
raise ValueError(
f"Unsupported file format: {self._dataset_path.suffix}. "
"Use .jsonl or .json file extensions."
)
data_frame = pd.DataFrame({self._data_column: data})
else:
raise ValueError("Dataset path or data must be provided")
assert data_frame is not None
self._data: pd.DataFrame = data_frame
if format and format not in ["oumi", "alpaca"]:
raise ValueError(
f"Invalid format: {format}. "
"Supported formats are 'oumi', and 'alpaca'."
)
self._format: str = format if format else self._detect_format(data_frame)
super().__init__(**kwargs)
def _detect_format(self, data_frame: pd.DataFrame) -> str:
"""Detect the format of the data based on the first item.
Args:
data_frame: The DataFrame containing the data.
Returns:
str: The detected format ("oumi", or "alpaca").
Raises:
ValueError: If the format cannot be detected.
"""
first_item = data_frame[self._data_column].iloc[0]
if not isinstance(first_item, dict):
raise ValueError(
"Invalid data format. Each item in the dataset should be a dictionary. "
f"Found type: {type(first_item)}. "
"Please check your data structure and try again."
)
if "messages" in first_item:
if isinstance(first_item["messages"], list) and all(
isinstance(m, dict) and "role" in m and "content" in m
for m in first_item["messages"]
):
return "oumi"
return "conversations"
elif all(key in first_item for key in ["instruction", "input", "output"]):
return "alpaca"
raise ValueError(
"Unable to auto-detect format. "
"The data structure doesn't match any supported format. "
"Please specify the format manually or ensure your data follows "
"one of these structures:\n"
"1. Conversations format: "
"{'messages': [{'role': 'user', 'content': '...'}, ...]}\n"
"2. Alpaca format: "
"{'instruction': '...', 'input': '...', 'output': '...'}\n"
)
@override
def _load_data(self) -> pd.DataFrame:
# Data is already loaded in __init__
return self._data
[docs]
@override
def transform_conversation(self, example: dict) -> Conversation:
"""Transform a single conversation example into a Conversation object.
Args:
example: The input example containing the messages or Alpaca-style turn.
Returns:
Conversation: A Conversation object containing the messages.
"""
conversation_dict = example[self._data_column]
if self._format == "oumi":
try:
return Conversation.model_validate(conversation_dict)
except Exception as e:
raise ValueError(
f"Invalid conversation format. "
f"Expected a dictionary with a 'messages' key "
f"containing a list of message dictionaries. Error: {str(e)}"
) from e
elif self._format == "alpaca":
return self._alpaca_to_conversation(conversation_dict)
else:
raise ValueError(f"Unsupported format: {self._format}")
def _alpaca_to_conversation(self, turn: dict) -> Conversation:
"""Convert an Alpaca-style turn to a Conversation object.
Args:
turn: A dictionary containing 'instruction', 'input', and 'output' keys.
Returns:
Conversation: A Conversation object representing the Alpaca-style turn.
Raises:
ValueError: If the turn doesn't contain the required keys.
"""
required_keys = ["instruction", "input", "output"]
if not all(key in turn for key in required_keys):
raise ValueError(
f"Invalid Alpaca format. The turn must contain all of these keys: "
"{required_keys}. "
f"Found keys: {list(turn.keys())}"
)
messages = [
{
"role": "user",
"content": f"{turn['instruction']}\n\n{turn['input']}".strip(),
},
{"role": "assistant", "content": turn["output"]},
]
return Conversation(messages=messages)
