2_collect_cluster_subset_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7b9df266-2541-4163-a57b-045f24d99b9a",
   "metadata": {},
   "source": [
    "## Get Relevant Paragraphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5c864b27-c45d-415c-b916-79489fda3520",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('strategyqa-data/strategyqa_dataset/strategyqa_train.json', 'r') as f:\n",
    "    qa = json.load(f)\n",
    "    \n",
    "question_set = set()\n",
    "with open('question_clusters/cluster4.json', 'r') as f:\n",
    "    loaded_strings = json.load(f)\n",
    "    \n",
    "for l in loaded_strings:\n",
    "    question_set.add(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "26b3d062-fcc0-4857-b38e-debf2ec8de4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "relevant_qas = []\n",
    "for q in qa:\n",
    "    if q['question'] in question_set:\n",
    "        relevant_qas.append(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6d2520ed-9077-44e7-a168-790386e05c1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections.abc import Iterable\n",
    "\n",
    "def flatten(xs):\n",
    "    for x in xs:\n",
    "        if isinstance(x, Iterable) and not isinstance(x, (str, bytes)):\n",
    "            yield from flatten(x)\n",
    "        else:\n",
    "            yield x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "3fb66158-ac11-464b-aad5-69eb026fabbf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 780 relevant paragraphs\n"
     ]
    }
   ],
   "source": [
    "relevant_paragraphs = set([])\n",
    "for q in relevant_qas:\n",
    "    relevant_paragraphs = relevant_paragraphs | set(flatten([q['evidence']]))\n",
    "    \n",
    "# two specific calls to strategyQA\n",
    "relevant_paragraphs.remove(\"no_evidence\")\n",
    "relevant_paragraphs.remove(\"operation\")\n",
    "    \n",
    "print(\"There are {} relevant paragraphs\".format(len(relevant_paragraphs)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c058b8d-d5ce-4c56-b3b6-237bd77dedac",
   "metadata": {},
   "source": [
    "## Total Number of Wikipedia Pages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "391da6f6-8e67-406f-899a-12de4fd60efb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Which come from 472 pages\n"
     ]
    }
   ],
   "source": [
    "removed_paragraph_endings = set([])\n",
    "for p in relevant_paragraphs:\n",
    "    removed_paragraph_endings.add(\"-\".join(p.split(\"-\")[:-1]))\n",
    "print(\"Which come from {} pages\".format(len(removed_paragraph_endings)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "148ba1a1-c4a6-4f58-a662-985296fabfdb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████| 36617357/36617357 [01:35<00:00, 385201.88it/s]\n"
     ]
    }
   ],
   "source": [
    "count_p = 0\n",
    "num_lines = 36617357\n",
    "with open('./strategyqa-data/corpus-enwiki-20200511-cirrussearch-parasv2.jsonl', 'r') as file:\n",
    "    for line in tqdm(file, total=num_lines):\n",
    "        obj = json.loads(line)\n",
    "\n",
    "        # Check if the 'title' parameter of the object exists in the given set\n",
    "        if 'title' in obj and obj['title'] in removed_paragraph_endings:\n",
    "            count_p += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "465f6656-868f-4cc2-8fa8-c8f8ac90d022",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overall there are 21396 related pages in total\n"
     ]
    }
   ],
   "source": [
    "print(\"Overall there are {} related pages in total\".format(count_p))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b7e2df-2374-4ffa-a2b9-af156ed09609",
   "metadata": {},
   "source": [
    "The 782 paragraphs will be used for training. Ideally, all 21k can be converted and stored in Qdrant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "fcfae6a7-444f-4219-84df-e0bd21b341a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('strategyqa-data/strategyqa_dataset/strategyqa_train_paragraphs.json', 'r') as f:\n",
    "    paragraph_content = json.load(f)\n",
    "\n",
    "paragraphs_to_index = []\n",
    "for p in relevant_paragraphs:\n",
    "    paragraphs_to_index.append(paragraph_content[p]['content'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "ed14ea8e-9c59-4d4b-bfee-40474b674e70",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('paragraph_content.json', 'w') as f:\n",
    "    json.dump(paragraphs_to_index, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama",
   "language": "python",
   "name": "llama"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}