augment_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/sumersao/miniconda3/envs/twitbits/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "import random\n",
    "import json\n",
    "\n",
    "random.seed(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Called by clean. Listify source documents, clean up summary (remove \"– \" at the beginning)\n",
    "\n",
    "input: single datapoint {'document': String, 'summary': String}\n",
    "output: {'document': List, 'summary': String}\n",
    "\"\"\"\n",
    "def clean_single_dp(datapoint):\n",
    "  docs_str = datapoint['document']\n",
    "  doc_list = docs_str.split(\"|||||\") #list of the source documents\n",
    "\n",
    "  sum = datapoint['summary']\n",
    "  summary_clean = sum[2:] #get rid of \"– \" at beginning of each summary\n",
    "\n",
    "  return doc_list, summary_clean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Clean each datapoint (listify docs, clean summaries) \n",
    "Create dicitonary where keys are number of source docs, in case we decide to \n",
    "aggregate our data based on #docs (to standardize input size)\n",
    "\n",
    "input: unaugmented multiNews dataset: List of {'document': String, 'summary': String}\n",
    "output: \n",
    "  all_data: List of {'document_list': List, 'clean_summary': String}\n",
    "  numdocs_dict: Dictionary where key = #source docs, value = list of datapoints - \n",
    "                each one is {'document': String, 'summary': String}\n",
    "\"\"\"\n",
    "def clean_orig(data): #takes all data\n",
    "  all_data = []\n",
    "\n",
    "  for point in data:\n",
    "    #augment single point\n",
    "    doc_list, summary_clean = clean_single_dp(point)\n",
    "\n",
    "    #add to all_data\n",
    "    all_data += [{'document_list': doc_list, 'clean_summary': summary_clean}]\n",
    "\n",
    "  return all_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Called by group_data. Aggregate given list of datapoints.\n",
    "\n",
    "input: list of datapoints, each of which is {document_list: [...], clean_summary: \"...\"}\n",
    "output: \n",
    "  combo_documents: list of doc lists\n",
    "  combo_summaries_str: String of combined summaries separated by \\n\\n\n",
    "  combo_summaries_list: list of summaries\n",
    "\"\"\"\n",
    "def combine_points(data_list):\n",
    "  combo_documents = [] #list of lists\n",
    "  combo_summaries_list = [] #list of strings\n",
    "\n",
    "  for point in data_list:\n",
    "    combo_documents += point['document_list']\n",
    "    combo_summaries_list += [point['clean_summary']]\n",
    "\n",
    "  combo_summaries_str = \"\\n\\n\".join(combo_summaries_list) #string (concatenated summaries separated by a black line)\n",
    "  # we also want to add a \"- \" to the beginning of each summary\n",
    "  combo_summaries_str = \"- \" + combo_summaries_str\n",
    "\n",
    "  return combo_documents, combo_summaries_str, combo_summaries_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Randomly partition dataset into groups of GROUP_SIZE, aggregate each group.\n",
    "\n",
    "input: full cleaned dataset - List of {'document_list': List, 'clean_summary': String}\n",
    "output: List of aggregated data {'documents': List, 'summary': String'\n",
    "\"\"\"\n",
    "def group_data(data):\n",
    "  GROUP_SIZE = 3\n",
    "\n",
    "  #partition\n",
    "  random.shuffle(data)\n",
    "  groups = [data[i:i+GROUP_SIZE] for i in range(0, len(data), GROUP_SIZE)]\n",
    "\n",
    "  #if last group is a lonely datapoint, merge it with the previous group\n",
    "  if len(groups[-1]) == 1:\n",
    "    groups[-2] += groups[-1]\n",
    "    groups = groups[:-1]\n",
    "\n",
    "  #aggregate each group\n",
    "  for i in range(len(groups)):\n",
    "    combo_docs, combo_sum_str, _ = combine_points(groups[i])\n",
    "    aggregated_group = {'documents': combo_docs, 'summary': combo_sum_str}\n",
    "    groups[i] = aggregated_group\n",
    "\n",
    "  return groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Compute distribution of #documents i.e. how many datapoints have x docs\n",
    "\n",
    "input: grouped and aggregated dataset - List of {'documents': List, 'summary': String'\n",
    "output: {int: int} key = number of documents, value = number of datapoints with that number of documents\n",
    "\"\"\"\n",
    "def get_group_numdoc_freq(grouped_data):\n",
    "    numdoc_frequency = {}\n",
    "\n",
    "    for data in grouped_data:\n",
    "        num_docs = len(data.get('documents', []))\n",
    "        if num_docs in numdoc_frequency:\n",
    "            numdoc_frequency[num_docs] += 1\n",
    "        else:\n",
    "            numdoc_frequency[num_docs] = 1\n",
    "\n",
    "    return numdoc_frequency\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "Full data augmentation process\n",
    "\"\"\"\n",
    "def augment_data(data):\n",
    "  clean_data = clean_orig(data)\n",
    "  augmented_data = group_data(clean_data)\n",
    "  return augmented_data #List of {'documents': List, 'summary': String}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset multi_news (/Users/sumersao/.cache/huggingface/datasets/multi_news/default/1.0.0/2f1f69a2bedc8ad1c5d8ae5148e4755ee7095f465c1c01ae8f85454342065a72)\n",
      "100%|██████████| 3/3 [00:00<00:00, 126.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train:  44972\n",
      "Val:  5622\n",
      "Test:  5622\n"
     ]
    }
   ],
   "source": [
    "dataset= load_dataset('multi_news')\n",
    "\n",
    "train = list(dataset['train'])\n",
    "val = list(dataset['validation'])\n",
    "test = list(dataset['test'])\n",
    "\n",
    "print(\"Train: \", len(train))\n",
    "print(\"Val: \", len(val))\n",
    "print(\"Test: \", len(test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Documents Frequency:\n",
      "#documents 9: 287\n",
      "#documents 7: 471\n",
      "#documents 6: 298\n",
      "#documents 8: 429\n",
      "#documents 10: 154\n",
      "#documents 11: 106\n",
      "#documents 5: 22\n",
      "#documents 14: 20\n",
      "#documents 13: 34\n",
      "#documents 12: 40\n",
      "#documents 15: 6\n",
      "#documents 19: 1\n",
      "#documents 16: 4\n",
      "#documents 18: 1\n",
      "#documents 17: 1\n"
     ]
    }
   ],
   "source": [
    "augmented_test = augment_data(test)\n",
    "numdoc_freq = get_group_numdoc_freq(augmented_test)\n",
    "\n",
    "print(\"Number of Documents Frequency:\")\n",
    "for numdocs, frequency in numdoc_freq.items():\n",
    "    print(f\"#documents {numdocs}: {frequency}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the augmented test set\n",
    "with open('augmented_test.json', 'w') as f:\n",
    "    json.dump(augmented_test, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "twitbits",
   "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}