{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import tiktoken\n",
    "import datasets\n",
    "import langdetect\n",
    "from semantic_text_splitter import TextSplitter\n",
    "from string import Template\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# * load dataset from jsonlines file\n",
    "dataset = datasets.load_dataset(\"json\", data_files=\"raw_data/pile/dedup-md5-pile-books3.jsonl\", split=\"train\")\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# * filter data by length\n",
    "enc = tiktoken.encoding_for_model(\"gpt-4\")\n",
    "\n",
    "def filter_length(examples):\n",
    "    res = []\n",
    "    for text in examples[\"text\"]:\n",
    "        token_len = len(enc.encode(text))\n",
    "        if token_len < 64_000:\n",
    "            res.append(False)\n",
    "        elif token_len > 80_000:\n",
    "            res.append(False)\n",
    "        else:\n",
    "            res.append(True)\n",
    "\n",
    "    return res\n",
    "\n",
    "\n",
    "dataset = dataset.filter(filter_length, batched=True, num_proc=32)\n",
    "\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# * filter non-English data\n",
    "dataset = dataset.filter(lambda x: langdetect.detect(x[\"text\"]) == \"en\", num_proc=32)\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# * make sure the data are not overlap\n",
    "used_dataset = datasets.load_dataset(\"json\", data_files=[\"backup_data/one_detail.book.jsonl\", \"backup_data/bio.book.jsonl\"], split=\"train\")\n",
    "\n",
    "dataset = dataset.filter(lambda x: x[\"md5\"] not in used_dataset[\"md5\"], num_proc=32)\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# * random sample\n",
    "dataset = dataset.train_test_split(test_size=150, seed=2024)[\"test\"]\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# * save data as the backup\n",
    "dataset.to_json(\"backup_data/multi_details.book.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasets.load_dataset(\"json\", data_files=\"backup_data/multi_details.book.jsonl\", split=\"train\")\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"Context information is below.\n",
    "---------------------\n",
    "${context}\n",
    "---------------------\n",
    "Given the context information and not prior knowledge.\n",
    "Generate content based on the below query.\n",
    "You are a Teacher/Professor. Your task is to setup 10 questions for an upcoming quiz/examination. The setup of questions should be divided into 2 steps. Step 1, you should find main characters and key plot in context information. Step 2, you should select a few related characters and plot and ask questions about them. The questions should meet following conditions. Condition 1, the expression of question should be diverse in nature. Condition 2, the questions should involve multiple related details to compare or reason. Restrict the questions to the context information provided.\n",
    "You must return the result in JSON: [{'question': <question>, 'answer': <answer>}, ..., {'question': <question>, 'answer': <answer>}]\"\"\"\n",
    "\n",
    "# * organize the data format\n",
    "jobs = []\n",
    "\n",
    "for idx, data in tqdm(enumerate(dataset)):\n",
    "    prompt = Template(template).substitute(context=data[\"text\"])\n",
    "    jobs.append({\n",
    "        \"model\": \"gpt-4-turbo-preview\", \n",
    "        \"temperature\": 0,\n",
    "        \"top_p\": 1.0,\n",
    "        \"max_tokens\": 4096,\n",
    "        \"messages\": [\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "        \"user\": f\"{idx}\",\n",
    "    })\n",
    "\n",
    "# * save, and then use Openai API script to generate data\n",
    "with open(\"data/multi_details.book.jsonl\", \"w\") as f:\n",
    "    for job in jobs:\n",
    "        json_string = json.dumps(job)\n",
    "        f.write(json_string + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
