{
 "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\", 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/bio.book.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasets.load_dataset(\"json\", data_files=\"backup_data/bio.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 Book Summarizer. Your task is to summarize the document. The task has 2 steps. In step 1, you should find the main characters. In step 2, you should summarize main characters' biography. The summary should be comprehensive and accurately reflect the main message.\n",
    "You must return the result in JSON: [{'character': <characters>, 'summary': <summary>}, ..., {'character': <characters>, 'summary': <summary>}]\"\"\"\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/bio.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
}
