BBH_Induct_split_preprocess_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "\n",
    "def split_jsonl_for_training(directory, seed=42, train_size=50):\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        for file in files:            \n",
    "            \n",
    "            if file.endswith(\".json\") and \"train\" not in file and \"test\" not in file:\n",
    "                file_path = os.path.join(root, file)\n",
    "            \n",
    "                with open(file_path, 'r', encoding='utf-8') as f:\n",
    "                    data = json.load(f) \n",
    "            \n",
    "                examples = data.get(\"examples\", [])\n",
    "\n",
    "                random.seed(seed)\n",
    "\n",
    "                if len(examples) > train_size:\n",
    "                    train_indices = random.sample(range(len(examples)), train_size)\n",
    "                else:\n",
    "                    train_indices = list(range(len(examples)))\n",
    "\n",
    "                train_data = [examples[i] for i in train_indices]\n",
    "                test_data = [examples[i] for i in range(len(examples)) if i not in train_indices]\n",
    "                \n",
    "                # print(train_data[0])\n",
    "                print(len(examples))\n",
    "\n",
    "                base_file_name = file.replace('.json', '')\n",
    "                train_file_path = os.path.join(root, f\"{base_file_name}.train.jsonl\")\n",
    "                test_file_path = os.path.join(root, f\"{base_file_name}.test.jsonl\")\n",
    "                \n",
    "                with open(train_file_path, 'w', encoding='utf-8') as f:\n",
    "                    for item in train_data:\n",
    "                        json_str = json.dumps(item) + \"\\n\"\n",
    "                        f.write(json_str)\n",
    "\n",
    "                with open(test_file_path, 'w', encoding='utf-8') as f:\n",
    "                    for item in test_data:\n",
    "                        json_str = json.dumps(item) + \"\\n\"\n",
    "                        f.write(json_str)\n",
    "\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "\n",
    "directory = \"BBH_preprocess/BBHI\" \n",
    "split_jsonl_for_training(directory)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7b3913313507bce4",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "def process_train_jsonl(directory, N_values=[3, 9]):\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        dir_name = os.path.basename(root) \n",
    "        for file in files:\n",
    "            if file.endswith(\".train.jsonl\"):\n",
    "                file_path = os.path.join(root, file)\n",
    "                base_file_name = file.replace('.train.jsonl', '')\n",
    "\n",
    "                system_content = None\n",
    "                is_content_same = True\n",
    "                user_contents = []\n",
    "\n",
    "                with open(file_path, 'r', encoding='utf-8') as f:\n",
    "                    for line_number, line in enumerate(f, 1):\n",
    "                        try:\n",
    "                            data = json.loads(line)\n",
    "                            if \"input\" in data:\n",
    "                                user_contents.append((data[\"input\"], data[\"target\"]))      \n",
    "                        except json.JSONDecodeError:\n",
    "                            print(f\"Json error in file {file_path} at line {line_number}.\")\n",
    "                        except KeyError as e:\n",
    "                            print(f\"Key error in file {file_path} at line {line_number}. Key: {e}\")\n",
    "                        if not is_content_same:\n",
    "                            break\n",
    "\n",
    "                \n",
    "                for N in N_values:\n",
    "                    examples_content = \"\"\n",
    "                    for i in range(min(N, len(user_contents))):\n",
    "                        input, target = user_contents[i]\n",
    "                        examples_content += f\"[Question]\\n{input}\\n\\n[Answer]\\n{target}\\n\\n\"\n",
    "                    with open(os.path.join(root, f\"{base_file_name}.{N}-shot_example.txt\"), 'w', encoding='utf-8') as f:\n",
    "                        f.write(examples_content)\n",
    "\n",
    "\n",
    "directory = \"BBH_preprocess/BBHI\"\n",
    "process_train_jsonl(directory)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e402eddda0bd2f0d",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "130dea52c4074236",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}