Hum-Works/lodestone-base-4096-v1

1{2 "cells": [3  {4   "cell_type": "markdown",5   "id": "e66bbb77-71f5-4d80-b766-f67144ea7a93",6   "metadata": {},7   "source": [8    "# Data Records\n",9    "\n",10    "## This notebook generates the data_records.json file where each entry in the resulting dictionary follows the form {filename: num_records} for every dataset we will use during training"11   ]12  },13  {14   "cell_type": "code",15   "execution_count": 39,16   "id": "74ad6613-44ff-435e-8550-df993e915677",17   "metadata": {18    "tags": []19   },20   "outputs": [],21   "source": [22    "# import relevant libraries\n",23    "import os\n",24    "import boto3\n",25    "import json\n",26    "from smart_open import open"27   ]28  },29  {30   "cell_type": "code",31   "execution_count": null,32   "id": "e2d53761-da0e-44f4-8a3e-1285bf810b03",33   "metadata": {34    "tags": []35   },36   "outputs": [],37   "source": [38    "s3 = boto3.resource('s3')\n",39    "my_bucket = s3.Bucket('lodestone-rnd')\n",40    "\n",41    "# collect all filenames from the data/ directory of the lodestone-rnd S3 bucket\n",42    "files = [\"\"]*((621+12+9+36)+1)\n",43    "for i, object_summary in enumerate(my_bucket.objects.filter(Prefix=\"data/\")):\n",44    "    files[i] = object_summary.key[5:]\n",45    "files = files[1:]\n",46    "files = [file for file in files if file != 'cnn_dailymail_splitted.json.gz']\n",47    "\n",48    "s3_client = boto3.client(\"s3\")\n",49    "\n",50    "# for each training dataset, store the number of records in a dictionary with the following form {filename: num_records}\n",51    "data_lengths = {}\n",52    "for file in files:\n",53    "    source_uri = f's3://lodestone-rnd/data/{file}'\n",54    "    # S2ORC_citations_abstracts.json.gz and amazon-qa.json.gz must be handled differently since each line in their training\n",55    "    # data is split into multiple records due to the fact that each query has multiple positive pair responses\n",56    "    if file in ['S2ORC_citations_abstracts.json.gz','amazon-qa.json.gz']:\n",57    "        length = 0\n",58    "        for json_line in open(source_uri, transport_params={\"client\": s3_client}):\n",59    "            data = json.loads(json_line.strip())\n",60    "            length += len(data['pos'])\n",61    "    else:\n",62    "        length = int(os.popen(f'aws s3 cp {source_uri} - | zcat | wc -l').read().rstrip())\n",63    "    data_lengths[f'{file}'] = length\n",64    "    \n",65    "# write the resulting dictionary to a .json file for future use during training\n",66    "with open('data_records.json', 'w') as fileout:\n",67    "    json.dump(data_lengths, fileout)"68   ]69  }70 ],71 "metadata": {72  "kernelspec": {73   "display_name": "conda_pytorch_p310",74   "language": "python",75   "name": "conda_pytorch_p310"76  },77  "language_info": {78   "codemirror_mode": {79    "name": "ipython",80    "version": 381   },82   "file_extension": ".py",83   "mimetype": "text/x-python",84   "name": "python",85   "nbconvert_exporter": "python",86   "pygments_lexer": "ipython3",87   "version": "3.10.10"88  }89 },90 "nbformat": 4,91 "nbformat_minor": 592}93