Parse a .json dataset file as one JSON document instead of line-by-line#7422
Parse a .json dataset file as one JSON document instead of line-by-line#7422vineethsaivs wants to merge 3 commits intounslothai:mainunslothai/unsloth:mainfrom vineethsaivs:fix/raw-text-json-document-parsingvineethsaivs/unsloth:fix/raw-text-json-document-parsingCopy head branch name to clipboard
Conversation
Both .json and .jsonl map to the "json_lines" reader, which parsed the file one line at a time with json.loads(line). A real .json file is a single JSON document (commonly a top-level list of records), so every line failed to parse, the whole document was dropped, and the reader returned an empty string; load_from_file then rejected the valid file with "is empty or contains only whitespace". Read the file and try json.loads on the whole content first (handling a list or a single object), falling back to line-by-line for true .jsonl. Add a CPU-only regression test.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d038c77f94
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if file_format == "plain_text" or file_format == "markdown": | ||
| return f.read() | ||
| elif file_format == "json_lines": | ||
| content = f.read() |
There was a problem hiding this comment.
Keep JSONL parsing streaming
Because .jsonl still maps to this json_lines handler, this f.read() now materializes the entire raw JSONL file before falling back to the per-line parser; the fallback also accumulates every decoded record before extracting text. Large training JSONL files that the previous loop processed one row at a time can now run out of memory, especially when rows contain large non-text fields. Please keep the .jsonl path streaming, or only use whole-document parsing for .json files.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, that was a real regression. Reading the whole file up front to try the document parse pulled a large .jsonl into memory, and the fallback then held every decoded record too, so the streaming loop the original code had was lost.
Fixed in d13a4b5: whole-document parsing is now only reached for a .json file, and .jsonl goes straight to a generator.
elif file_format == "json_lines":
if Path(file_path).suffix.lower() == ".json":
# A .json file is a single JSON document (commonly a list
# of records), so parsing it per line drops the whole file.
try:
parsed = json.load(f)
records = parsed if isinstance(parsed, list) else [parsed]
except json.JSONDecodeError:
# Some files carry JSON Lines under a .json name.
f.seek(0)
records = self._iter_json_lines(f)
else:
# A .jsonl file is one JSON value per line: stay streaming so
# a large file is never held in memory all at once.
records = self._iter_json_lines(f)
lines = []
for data in records:
text = self._extract_text_from_json(data)
if text:
lines.append(text)
return "\n\n".join(lines)_iter_json_lines yields one record at a time (skipping blank and malformed lines, same as the original loop), so on the .jsonl path nothing more than a single line is resident, and the records list no longer exists at all. json.load(f) on the .json path reads the document whole, which is unavoidable for a single JSON value, and that is the file shape this PR is about.
I kept the per-line fallback for .json, since files carrying JSON Lines under a .json name are common and they parsed fine before this PR. It only runs after a whole-document parse has already failed on a .json file, so it never costs the .jsonl path anything.
Added test_jsonl_is_never_materialized, which wraps the handle so read() and seek() raise, and it fails on the previous commit:
E AssertionError: .jsonl was read whole instead of streamed line by line
Also added test_json_holding_json_lines_still_falls_back to pin the mislabelled-.json path. All 4 tests pass on the new commit.
The previous commit read the whole file before trying a whole-document parse, so a large .jsonl file that used to be processed one row at a time was materialised in memory, along with every decoded record. Whole-document parsing is only needed for a .json file, so branch on the extension: .json parses with json.load (with the per-line fallback kept for files that carry JSON Lines under a .json name), and .jsonl goes straight to a generator that yields one record per line and never holds the file.
Summary
RawTextDataLoadermaps both.jsonand.jsonlto the same"json_lines"reader (SUPPORTED_FORMATS), and that reader parses the file one line at a time:A real
.jsonfile is a single JSON document, usually a top-level list of records (the common instruct-dataset shape) or a pretty-printed object spanning many lines. None of its individual lines are valid standalone JSON, so everyjson.loads(line)raises, all lines are skipped, and the reader returns"".load_from_filethen rejects the perfectly valid file withValueError: File '...' is empty or contains only whitespace. Only true line-delimited.jsonlworked.Fix
Read the file and try
json.loadson the whole content first: a list stays a list of records, a single object becomes one record. Fall back to the existing line-by-line parse for true.jsonl(a multi-line JSONL file fails the whole-file parse and drops through). Text extraction and the.jsonlpath are otherwise unchanged.Test
tests/test_raw_text_json_loading.pystubsdatasetsand execsraw_text.pydirectly (noimport unsloth, so no GPU / unsloth_zoo), then checks that a.jsonarray file and a.jsonlfile both read back their records. The.jsoncase fails before this change (the reader returns"") and passes after; the.jsonlcase is a regression guard. Wired into the Bucket-A CPU CI step.