CoolFace
Datasetpublic

aluncstokes/mathpile_arxiv_subset_tiny

MathPile ArXiv (subset) Description This dataset consists of a toy subset of 8834 (5000 training + 3834 testing) TeX files found in the arXiv subset of MathPile, used for testing. You should not use this dataset. Training and testing sets are already split Source The data was obtained from the training + validation portion of the arXiv subset of MathPile. Format Given as JSONL files of JSON dicts each containing the single key:… See the full description on the dataset page: https://huggingface.co/datasets/aluncstokes/mathpile_arxiv_subset_tiny.

sourceHugging Faceupdated 3y agoView on Hugging Face
0likes50downloads
process_documents.ipynb268 linesDownload Raw Back to root
1{2 "cells": [3  {4   "cell_type": "code",5   "execution_count": 1,6   "metadata": {},7   "outputs": [8    {9     "name": "stdout",10     "output_type": "stream",11     "text": [12      "Note: you may need to restart the kernel to use updated packages.\n",13      "Note: you may need to restart the kernel to use updated packages.\n",14      "Note: you may need to restart the kernel to use updated packages.\n",15      "Note: you may need to restart the kernel to use updated packages.\n",16      "Note: you may need to restart the kernel to use updated packages.\n"17     ]18    }19   ],20   "source": [21    "%pip install -q -U tiktoken\n",22    "%pip install -q -U langchain\n",23    "%pip install -q -U langdetect\n",24    "%pip install -q -U git+https://github.com/huggingface/transformers\n",25    "%pip install -q -U sentencepiece\n",26    "%pip install -q -U protobuf\n",27    "%pip install -q -U tqdm"28   ]29  },30  {31   "cell_type": "code",32   "execution_count": 2,33   "metadata": {},34   "outputs": [],35   "source": [36    "import os\n",37    "import json\n",38    "import warnings\n",39    "from tqdm import tqdm\n",40    "\n",41    "# import tiktoken\n",42    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",43    "from langdetect import detect, detect_langs\n",44    "from transformers import AutoTokenizer\n",45    "\n",46    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"True\""47   ]48  },49  {50   "cell_type": "code",51   "execution_count": 12,52   "metadata": {},53   "outputs": [],54   "source": [55    "tokeniser = AutoTokenizer.from_pretrained(\n",56    "    \"mistralai/Mistral-7B-v0.1\", use_fast=True)\n",57    "\n",58    "\n",59    "def string_token_length(text):\n",60    "    return len(tokeniser(text, add_special_tokens=False).input_ids)\n",61    "\n",62    "\n",63    "def write_jsonl(data, filename):\n",64    "    with open(filename, \"w\") as f:\n",65    "        for entry in data:\n",66    "            f.write(json.dumps(entry) + \"\\n\")\n",67    "\n",68    "\n",69    "def chunkify_tex(tex_text, chunk_size=4086, chunk_overlap=1536):\n",70    "    splitter = RecursiveCharacterTextSplitter(\n",71    "        [\n",72    "            r\"(?<=\\n)(?=\\\\section{)\",\n",73    "            r\"(?<=\\n)(?=\\\\subsection{)\",\n",74    "            r\"(?<=\\n)(?=\\\\subsubsection{)\",\n",75    "            r\"(?<=\\\\end{proof})\",\n",76    "            r\"(?<=\\\\qed)\",\n",77    "            r\"\\n\\n\\n\",\n",78    "            r\"\\n\\n\",\n",79    "        ],\n",80    "        keep_separator=True,\n",81    "        is_separator_regex=True,\n",82    "        chunk_size=chunk_size,\n",83    "        chunk_overlap=chunk_overlap,\n",84    "        length_function=string_token_length,\n",85    "    )\n",86    "\n",87    "    bits = splitter.split_text(tex_text)\n",88    "\n",89    "    return bits"90   ]91  },92  {93   "cell_type": "code",94   "execution_count": 4,95   "metadata": {},96   "outputs": [],97   "source": [98    "# Metadata object\n",99    "# {\n",100    "#   \"set\": str,\n",101    "#   \"id\": int,\n",102    "#   \"char_len\": int,\n",103    "#   \"tok_len\": int,\n",104    "#   \"lang\": str,\n",105    "#   \"text\": str\n",106    "#   \"num_chunks\": int\n",107    "#   \"chunk_lengths\": List[int]\n",108    "#   \"chunks\": List[str]\n",109    "# }\n",110    "\n",111    "with open(\"mathpile_arxiv_subset_tiny/train.jsonl\", \"r\") as f:\n",112    "    train = [json.loads(l) for l in f.readlines()]\n",113    "\n",114    "with open(\"mathpile_arxiv_subset_tiny/test.jsonl\", \"r\") as f:\n",115    "    test = [json.loads(l) for l in f.readlines()]"116   ]117  },118  {119   "cell_type": "code",120   "execution_count": 5,121   "metadata": {},122   "outputs": [123    {124     "name": "stderr",125     "output_type": "stream",126     "text": [127      "8438it [02:30, 56.19it/s]\n"128     ]129    }130   ],131   "source": [132    "data = [\n",133    "    {\n",134    "        \"set\": \"train\",\n",135    "        \"id\": f'0.{j}',\n",136    "        \"char_len\": len(entry[\"text\"]),\n",137    "        # \"tok_len\": string_token_length(entry[\"text\"]),\n",138    "        \"lang\": detect(entry[\"text\"]),\n",139    "        \"text\": entry[\"text\"],\n",140    "    }\n",141    "    for j, entry in tqdm(enumerate(train), total=len(train))\n",142    "] + [\n",143    "    {\n",144    "        \"set\": \"test\",\n",145    "        \"id\": f'1.{j}',\n",146    "        \"char_len\": len(entry[\"text\"]),\n",147    "        # \"tok_len\": string_token_length(entry[\"text\"]),\n",148    "        \"lang\": detect(entry[\"text\"]),\n",149    "        \"text\": entry[\"text\"],\n",150    "    }\n",151    "    for j, entry in tqdm(enumerate(test), total=len(test))\n",152    "]"153   ]154  },155  {156   "cell_type": "code",157   "execution_count": 13,158   "metadata": {},159   "outputs": [160    {161     "name": "stderr",162     "output_type": "stream",163     "text": [164      "100%|██████████| 8438/8438 [15:17<00:00,  9.20it/s]\n"165     ]166    }167   ],168   "source": [169    "for j, datum in tqdm(enumerate(data), total=len(data)):\n",170    "    text = datum[\"text\"]\n",171    "    chunks = chunkify_tex(text, chunk_size=4095, chunk_overlap=1536)\n",172    "    chunk_lengths = [string_token_length(chunk) for chunk in chunks]\n",173    "    data[j][\"num_chunks\"] = len(chunks)\n",174    "    data[j][\"chunk_lengths\"] = chunk_lengths\n",175    "    data[j][\"chunks\"] = chunks\n",176    "    data[j][\"tok_len\"] = sum(chunk_lengths)"177   ]178  },179  {180   "cell_type": "code",181   "execution_count": 18,182   "metadata": {},183   "outputs": [],184   "source": [185    "data_train = [datum for datum in data if datum[\"set\"] == \"train\"]\n",186    "data_test = [datum for datum in data if datum[\"set\"] == \"test\"]\n",187    "\n",188    "chunked_train = []\n",189    "chunked_test = []\n",190    "\n",191    "for datum in data_train:\n",192    "    total_document_tokens = datum[\"tok_len\"]\n",193    "    document_id = datum[\"id\"]\n",194    "    for j in range(len(datum[\"chunks\"])):\n",195    "        chunk_id = f\"{document_id}.{j}\"\n",196    "        chunk_text = datum[\"chunks\"][j]\n",197    "        chunk_num_tokens = datum[\"chunk_lengths\"][j]\n",198    "        chunked_train.append(\n",199    "            {\n",200    "                \"set\": \"train\",\n",201    "                \"id\": chunk_id,\n",202    "                \"chunk_text\": chunk_text,\n",203    "                \"chunk_num_tokens\": chunk_num_tokens,\n",204    "                \"document_num_tokens\": total_document_tokens,\n",205    "                \"document_language\": datum[\"lang\"],\n",206    "            }\n",207    "        )\n",208    "\n",209    "for datum in data_test:\n",210    "    total_document_tokens = datum[\"tok_len\"]\n",211    "    document_id = datum[\"id\"]\n",212    "    for j in range(len(datum[\"chunks\"])):\n",213    "        chunk_id = f\"{document_id}.{j}\"\n",214    "        chunk_text = datum[\"chunks\"][j]\n",215    "        chunk_num_tokens = datum[\"chunk_lengths\"][j]\n",216    "        chunked_test.append(\n",217    "            {\n",218    "                \"set\": \"test\",\n",219    "                \"id\": chunk_id,\n",220    "                \"chunk_text\": chunk_text,\n",221    "                \"chunk_num_tokens\": chunk_num_tokens,\n",222    "                \"document_num_tokens\": total_document_tokens,\n",223    "                \"document_language\": datum[\"lang\"],\n",224    "            }\n",225    "        )"226   ]227  },228  {229   "cell_type": "code",230   "execution_count": 21,231   "metadata": {},232   "outputs": [],233   "source": [234    "write_jsonl(chunked_train, \"mathpile_arxiv_subset_tiny/train_chunked.jsonl\")\n",235    "write_jsonl(chunked_test, \"mathpile_arxiv_subset_tiny/test_chunked.jsonl\")"236   ]237  },238  {239   "cell_type": "code",240   "execution_count": null,241   "metadata": {},242   "outputs": [],243   "source": []244  }245 ],246 "metadata": {247  "kernelspec": {248   "display_name": "venv",249   "language": "python",250   "name": "python3"251  },252  "language_info": {253   "codemirror_mode": {254    "name": "ipython",255    "version": 3256   },257   "file_extension": ".py",258   "mimetype": "text/x-python",259   "name": "python",260   "nbconvert_exporter": "python",261   "pygments_lexer": "ipython3",262   "version": "3.9.6"263  }264 },265 "nbformat": 4,266 "nbformat_minor": 2267}268