{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# The jsonl loading challenge" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "N = 10_000_000" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Make the test data" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000000/10000000 [00:09<00:00, 1031209.45it/s]\n" ] } ], "source": [ "import json\n", "from tqdm import tqdm\n", "\n", "line = {\n", " \"id\": 0,\n", " \"text\": \"This is a test record 0\",\n", " \"tags\": [\"jazz\", \"funk\", \"soul\"],\n", " \"audio\": {\n", " \"transcript\": \"This is a test transcript 0\" * 50,\n", " \"audio_path\": \"This is a test audio path 0\",\n", " \"audio_url\": \"This is a test audio url 0\",\n", " \"audio_duration\": 1.0,\n", " \"audio_sample_rate\": 16000,\n", " \"energy\": [0.2] * 60,\n", " },\n", "}\n", "line_s = json.dumps(line)\n", "with open(\"/tmp/test_data.jsonl\", \"w\") as f:\n", " for i in tqdm(range(N)):\n", " f.write(line_s + \"\\n\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Example of how to load the data" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "-rw-r--r--@ 1 nat wheel 18G Mar 11 11:04 /tmp/test_data.jsonl\n" ] } ], "source": [ "!ls -lh /tmp/test_data.jsonl" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Requirement already satisfied: orjson in /home/victor/anaconda3/envs/suno/lib/python3.10/site-packages (3.9.1)\n", "\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.1.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n" ] } ], "source": [ "# !pip install orjson # faster json loading" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Exception ignored in: >\n", "Traceback (most recent call last):\n", " File \"/Users/nat/Repos/nat/json_rows/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py\", line 775, in _clean_thread_parent_frames\n", " def _clean_thread_parent_frames(\n", "\n", "KeyboardInterrupt: \n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "{'id': 0, 'text': 'This is a test record 0', 'tags': ['jazz', 'funk', 'soul'], 'audio': {'transcript': 'This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0This is a test transcript 0', 'audio_path': 'This is a test audio path 0', 'audio_url': 'This is a test audio url 0', 'audio_duration': 1.0, 'audio_sample_rate': 16000, 'energy': [0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]}}\n", "Time taken: 137.4187572002411\n" ] } ], "source": [ "import orjson\n", "import mmap\n", "import time\n", "\n", "# this takes 2:07 minutes on the cluster\n", "\n", "path = \"/tmp/test_data.jsonl\"\n", "\n", "metadata = [None] * N\n", "start = time.time()\n", "with open(path, \"r\") as f:\n", " mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)\n", " i = 0\n", "\n", " while True:\n", " line = mm.readline()\n", " if line == b\"\":\n", " break\n", " data = orjson.loads(line)\n", " metadata[i] = data\n", " i += 1\n", "\n", " mm.close()\n", "print(metadata[-1])\n", "end = time.time() - start\n", "print(f\"Time taken: {end}\")\n" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "assert len(metadata) == N\n", "assert metadata[-1][\"audio\"][\"transcript\"] == \"This is a test transcript 0\" * 50" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 2 }