post_process.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import glob\n",
    "from collections import defaultdict\n",
    "from tqdm import tqdm\n",
    "import re\n",
    "\n",
    "def post_process(model_name, extract_func):\n",
    "    # Directory where your files are stored\n",
    "    directory = './results/'\n",
    "\n",
    "    # Pattern to match the files of interest\n",
    "    pattern = directory + f'*{model_name}*.json'\n",
    "\n",
    "    # Process each file\n",
    "    for file_path in glob.glob(pattern):\n",
    "        print(f\"Processing {file_path}...\")\n",
    "        \n",
    "        # Read the file\n",
    "        with open(file_path, 'r') as f:\n",
    "            data = json.load(f)\n",
    "        \n",
    "        # Extract triples\n",
    "        text_triples = extract_func(data)\n",
    "        \n",
    "        # Define new file name for processed results\n",
    "        base_name = os.path.basename(file_path)\n",
    "        new_file_name = base_name.replace('results', 'processed_results')\n",
    "        new_file_path = os.path.join('./processed_results/', new_file_name)\n",
    "        \n",
    "        # Write the processed data to a new file\n",
    "        with open(new_file_path, 'w') as f:\n",
    "            json.dump(text_triples, f, indent=6)\n",
    "        \n",
    "        print(f\"Processed data saved to {new_file_path}\")\n",
    "        print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GPT Family: gpt-3.5-turbo-instruct & gpt-4 & text-davinci-003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "from collections import defaultdict\n",
    "from tqdm import tqdm\n",
    "\n",
    "def extract_triples_gpt(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        # Use regex to find all JSON array-like structures\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/nyt10m_rand_500_gpt-3.5-turbo-instruct_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 28885.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_gpt-3.5-turbo-instruct_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_gpt-3.5-turbo-1106_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 45519.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_gpt-3.5-turbo-1106_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_gpt-3.5-turbo-1106_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 58369.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_gpt-3.5-turbo-1106_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_gpt-3.5-turbo-instruct_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 10013.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_gpt-3.5-turbo-instruct_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_gpt-3.5-turbo-instruct_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 10366.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_gpt-3.5-turbo-instruct_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_gpt-3.5-turbo-1106_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 13554.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_gpt-3.5-turbo-1106_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_gpt-3.5-turbo-instruct_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 33845.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_gpt-3.5-turbo-instruct_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_gpt-3.5-turbo-instruct_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 38253.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_gpt-3.5-turbo-instruct_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_gpt-3.5-turbo-1106_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 12046.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_gpt-3.5-turbo-1106_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_gpt-3.5-turbo-1106_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 14472.35it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_gpt-3.5-turbo-1106_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-3.5-turbo-instruct_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 36115.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-3.5-turbo-instruct_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-3.5-turbo-1106_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 63563.54it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-3.5-turbo-1106_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-3.5-turbo-1106_54.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 61307.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-3.5-turbo-1106_54.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-3.5-turbo-1106_64.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 61706.35it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-3.5-turbo-1106_64.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-3.5-turbo-1106_74.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 66582.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-3.5-turbo-1106_74.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-3.5-turbo-1106_84.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 63551.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-3.5-turbo-1106_84.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_gpt-4_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 34730.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_gpt-4_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 8192.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_gpt-4_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 10205.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_gpt-4_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_gpt-4_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 12365.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_gpt-4_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 33932.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_gpt-4_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 43681.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_gpt-4_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_gpt-4_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 40340.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_gpt-4_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 9860.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 39754.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 45459.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 39927.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 40591.35it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_54.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 40618.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_54.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_64.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 40126.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_64.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_74.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 38863.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_74.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_84.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 39366.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_84.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_text-davinci-003_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 36684.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_text-davinci-003_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_text-davinci-003_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 11625.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_text-davinci-003_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_text-davinci-003_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 14048.68it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_text-davinci-003_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_text-davinci-003_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 40774.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_text-davinci-003_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_text-davinci-003_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 40659.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_text-davinci-003_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_text-davinci-003_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 11605.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_text-davinci-003_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 8560.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 32263.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 9831.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 41324.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 38924.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_gpt-4-1106-preview_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 42454.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_gpt-4-1106-preview_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_54.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 40142.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_54.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_64.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 39713.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_64.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_74.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 40241.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_74.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_gpt-4-1106-preview_84.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 40198.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_gpt-4-1106-preview_84.json\n",
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('gpt-3.5-turbo', extract_triples_gpt)\n",
    "post_process('gpt-4', extract_triples_gpt)\n",
    "post_process('davinci', extract_triples_gpt)\n",
    "post_process('gpt-4-1106', extract_triples_gpt)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LLAMA Family: Vicuna-7B, Vicuna-33B, LLAMA-2-7B, LLAMA-2-70B, Wizard-70B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_triples_vicuna(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        triples_str = triples_str.split(\"\\nASSISTANT:\")[1]\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/tacred_rand_800_vicuna-1.5-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 51329.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_vicuna-1.5-7b_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_vicuna-1.3-33b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 33163.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_vicuna-1.3-33b_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_vicuna-1.3-33b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 14583.68it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_vicuna-1.3-33b_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_vicuna-1.3-33b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 15844.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_vicuna-1.3-33b_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_vicuna-1.3-33b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 10099.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_vicuna-1.3-33b_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_vicuna-1.5-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 6288.41it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_vicuna-1.5-7b_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_vicuna-1.3-33b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 5486.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_vicuna-1.3-33b_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_vicuna-1.5-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 5721.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_vicuna-1.5-7b_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_vicuna-1.5-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 28522.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_vicuna-1.5-7b_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_vicuna-1.5-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 18707.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_vicuna-1.5-7b_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_vicuna-1.3-33b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 11921.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_vicuna-1.3-33b_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_vicuna-1.5-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 24444.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_vicuna-1.5-7b_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_vicuna-1.3-33b_54.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 27540.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_vicuna-1.3-33b_54.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_vicuna-1.3-33b_64.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 27852.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_vicuna-1.3-33b_64.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_vicuna-1.3-33b_74.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 28340.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_vicuna-1.3-33b_74.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_vicuna-1.3-33b_84.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 28199.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_vicuna-1.3-33b_84.json\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('vicuna', extract_triples_vicuna)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LLAMA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_triples_llama(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        triples_str = triples_str.split(\"[/INST]\")[1]\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/tacred_rand_800_llama-2-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 14228.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_llama-2-70b_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_llama-2-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 5196.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_llama-2-70b_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_llama-2-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 15812.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_llama-2-70b_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_llama-2-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 27320.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_llama-2-70b_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_llama-2-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 32192.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_llama-2-7b_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_llama-2-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 10004.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_llama-2-70b_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_llama-2-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 5128.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_llama-2-7b_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_llama-2-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 14341.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_llama-2-70b_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_llama-2-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 24251.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_llama-2-7b_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_llama-2-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 23118.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_llama-2-7b_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_llama-2-7b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 37186.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_llama-2-7b_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_llama-2-7b_44.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 4707.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_llama-2-7b_44.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_llama-2-70b_54_.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 22739.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_llama-2-70b_54_.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_llama-2-70b_54.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 15498.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_llama-2-70b_54.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_llama-2-70b_74.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 25576.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_llama-2-70b_74.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_llama-2-70b_64.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 30523.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_llama-2-70b_64.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_llama-2-70b_84.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 24695.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_llama-2-70b_84.json\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('llama', extract_triples_llama)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wizard-70B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/docred_rand_200_wizardlm-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 20326.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_wizardlm-70b_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_wizardlm-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 33787.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_wizardlm-70b_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_wizardlm-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 31073.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_wizardlm-70b_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_wizardlm-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 34864.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_wizardlm-70b_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_wizardlm-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 20228.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_wizardlm-70b_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_wizardlm-70b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 58745.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_wizardlm-70b_1.json\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('wizard', extract_triples_vicuna)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mistral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_triples_mistral(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        triples_str = triples_str.split(\"[/INST]\")[1]\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/nyt10m_rand_500_mistral_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 31093.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_mistral_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_mistral_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 10405.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_mistral_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_mistral_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 12486.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_mistral_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_mistral_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 42496.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_mistral_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_mistral_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 60483.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_mistral_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_mistral_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 67202.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_mistral_1.json\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('mistral', extract_triples_mistral)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GALACTICA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_triples_galactica(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/wiki80_rand_800_galactica-30b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 31276.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_galactica-30b_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_galactica-30b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 23826.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_galactica-30b_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_galactica-30b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 30907.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_galactica-30b_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_galactica-30b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 579724.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_galactica-30b_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_galactica-30b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 283974.54it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_galactica-30b_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_galactica-30b_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 23600.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_galactica-30b_1.json\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('galactica', extract_triples_galactica)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### zephyr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_triples_zephyr(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        triples_str = triples_str.split(\"\\n<|assistant|>\\n\")[1]\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/tacred_rand_800_zephyr-7b-beta_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 28055.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_zephyr-7b-beta_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_zephyr-7b-beta_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 9053.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_zephyr-7b-beta_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_zephyr-7b-beta_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 6556.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_zephyr-7b-beta_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_zephyr-7b-beta_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 21434.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_zephyr-7b-beta_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_zephyr-7b-beta_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 29816.20it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_zephyr-7b-beta_1.json\n",
      "\n",
      "Processing ./results/wiki80_rand_800_zephyr-7b-beta_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 800/800 [00:00<00:00, 35002.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_zephyr-7b-beta_1.json\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('zephyr', extract_triples_zephyr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### openchat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_triples_openchat(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        triples_str = triples_str.split(\"GPT4 Correct Assistant:\")[1]\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/wiki80_rand_800_openchat_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 41204.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki80_rand_800_openchat_1.json\n",
      "\n",
      "Processing ./results/tacred_rand_800_openchat_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 800/800 [00:00<00:00, 14167.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/tacred_rand_800_openchat_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_openchat_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 33160.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_openchat_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_openchat_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 15682.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_openchat_1.json\n",
      "\n",
      "Processing ./results/docred_rand_200_openchat_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 11230.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/docred_rand_200_openchat_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_openchat_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 40292.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_openchat_1.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_openchat_54.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 40652.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_openchat_54.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_openchat_64.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 39271.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_openchat_64.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_openchat_74.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 38385.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_openchat_74.json\n",
      "\n",
      "Processing ./results/wiki20m_rand_500_openchat_84.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 39358.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/wiki20m_rand_500_openchat_84.json\n",
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('openchat', extract_triples_openchat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 207,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "from collections import defaultdict\n",
    "\n",
    "def extract_triples(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in data.items():\n",
    "        # triples_str = triples_str.split(\"\\nASSISTANT:\")[1] (For Vicuna)\n",
    "        # triples_str = triples_str.split(\"[/INST]\")[1] (For LLaMA & Mistral)\n",
    "        # triples_str = triples_str.split(\"GPT4 Correct Assistant:\")[1] (For OpenChat)\n",
    "        repeat_map = {}\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            triple = json.loads(match)\n",
    "            if len(triple) == 3:\n",
    "                triple_list.append(triple)\n",
    "            if str(triple) not in repeat_map:\n",
    "                repeat_map[str(triple)] = 1\n",
    "            else:\n",
    "                repeat_map[str(triple)] += 1\n",
    "\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            triples = json.loads(match)\n",
    "            for triple in triples:\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                elif repeat_map[str(triple)] == 1:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "                    continue\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "\n",
    "        text_triples[text] = triple_list\n",
    "    return text_triples\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## semi-open"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "from collections import defaultdict\n",
    "from tqdm import tqdm\n",
    "\n",
    "def extract_triples_semi(data):\n",
    "    text_triples = defaultdict(list)\n",
    "    for text, triples_str in tqdm(data.items()):\n",
    "        # Use regex to find all JSON array-like structures\n",
    "        repeat_map = {}\n",
    "        triples_str = triples_str.replace(\"\\'\", '\\\"')\n",
    "        matches_1 = re.findall(r'\\[.*?\\]', triples_str, re.DOTALL)\n",
    "        matches_2 = re.findall(r'\\[\\[.*?\\]\\]', triples_str, re.DOTALL)\n",
    "        triple_list = []\n",
    "        for match in matches_1:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triple = json.loads(match)\n",
    "                if len(triple) == 3:\n",
    "                    triple_list.append(triple)\n",
    "                if str(triple) not in repeat_map:\n",
    "                    repeat_map[str(triple)] = 1\n",
    "                else:\n",
    "                    repeat_map[str(triple)] += 1\n",
    "\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                    continue\n",
    "        text_triples[text] = triple_list\n",
    "        \n",
    "        \n",
    "        for match in matches_2:\n",
    "            try:\n",
    "                # Attempt to parse each match as JSON\n",
    "                triples = json.loads(match)\n",
    "                for triple in triples:\n",
    "                    if str(triple) not in repeat_map:\n",
    "                        repeat_map[str(triple)] = 1\n",
    "                    elif repeat_map[str(triple)] == 1:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                        continue\n",
    "                    else:\n",
    "                        repeat_map[str(triple)] += 1\n",
    "                    if len(triple) == 3:\n",
    "                        triple_list.append(triple)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the match is not valid JSON\n",
    "                continue\n",
    "            \n",
    "        triple_list_ = []\n",
    "        for triple in triple_list:\n",
    "            subject_, predicate_, object_ = triple\n",
    "            if \":\" in subject_:\n",
    "                subject_ = subject_.split(\":\")[0]\n",
    "            if \":\" in object_:\n",
    "                object_ = object_.split(\":\")[0]\n",
    "            triple = (subject_, predicate_, object_)\n",
    "            triple_list_.append(triple)\n",
    "                \n",
    "        text_triples[text] = triple_list_\n",
    "    return text_triples\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ./results/nyt10m_gpt-3.5_semi_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00, 9279.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_gpt-3.5_semi_1.json\n",
      "\n",
      "Processing ./results/nyt10m_rand_500_gpt-3.5_semi_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:00<00:00, 30943.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/nyt10m_rand_500_gpt-3.5_semi_1.json\n",
      "\n",
      "Processing ./results/cdr_rand_200_gpt-3.5_semi_1.json...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 22504.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed data saved to ./processed_results/cdr_rand_200_gpt-3.5_semi_1.json\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "post_process('semi', extract_triples_semi)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## closed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('/home/pj20/GREScore/results/nyt10m_gpt-3.5_closed_1.json', 'r') as f:\n",
    "    data = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "results = defaultdict(list)\n",
    "\n",
    "for text, triples in data.items():\n",
    "    for triple in triples:\n",
    "        subject_, relation, object_ = triple\n",
    "        if relation[0] == ' ':\n",
    "            relation = relation[1:]\n",
    "        results[text].append([subject_, relation, object_])\n",
    "        \n",
    "with open('/home/pj20/GREScore/processed_results/nyt10m_gpt-3.5_closed_1.json', 'w') as f:\n",
    "    json.dump(results, f, indent=6)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.9 ('textgen')",
   "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.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "72e4df05eb8a5f0bdb80d33a9292878c0e2ee2d3f5bd7213be092f7893c82c13"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}