gen_test_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "46b362f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./test_data/Counting_Stars_EN_multi-evidence-retrieval-searching_128000_32_32.jsonl\n",
      "撒了32次星星\n",
      "18866\n",
      "撒了32次星星\n",
      "37223\n",
      "撒了32次星星\n",
      "55676\n",
      "撒了32次星星\n",
      "72449\n",
      "撒了32次星星\n",
      "90226\n",
      "撒了32次星星\n",
      "108808\n",
      "撒了32次星星\n",
      "126043\n",
      "撒了32次星星\n",
      "143722\n",
      "撒了32次星星\n",
      "161214\n",
      "撒了32次星星\n",
      "178738\n",
      "撒了32次星星\n",
      "196217\n",
      "撒了32次星星\n",
      "214366\n",
      "撒了32次星星\n",
      "232257\n",
      "撒了32次星星\n",
      "250212\n",
      "撒了32次星星\n",
      "267685\n",
      "撒了32次星星\n",
      "284783\n",
      "撒了32次星星\n",
      "302384\n",
      "撒了32次星星\n",
      "320187\n",
      "撒了32次星星\n",
      "338073\n",
      "撒了32次星星\n",
      "355557\n",
      "撒了32次星星\n",
      "374231\n",
      "撒了32次星星\n",
      "392587\n",
      "撒了32次星星\n",
      "410318\n",
      "撒了32次星星\n",
      "428433\n",
      "撒了32次星星\n",
      "446051\n",
      "撒了32次星星\n",
      "463849\n",
      "撒了32次星星\n",
      "482058\n",
      "撒了32次星星\n",
      "499498\n",
      "撒了32次星星\n",
      "516604\n",
      "撒了32次星星\n",
      "533982\n",
      "撒了32次星星\n",
      "552214\n",
      "撒了32次星星\n",
      "570151\n",
      "共计32条数据\n",
      "./test_data/Counting_Stars_ZH_multi-evidence-retrieval-searching_128000_32_32.jsonl\n",
      "撒了32次星星\n",
      "3289\n",
      "撒了32次星星\n",
      "6089\n",
      "撒了32次星星\n",
      "8889\n",
      "撒了32次星星\n",
      "11689\n",
      "撒了32次星星\n",
      "14489\n",
      "撒了32次星星\n",
      "17289\n",
      "撒了32次星星\n",
      "20089\n",
      "撒了32次星星\n",
      "22889\n",
      "撒了32次星星\n",
      "25689\n",
      "撒了32次星星\n",
      "28489\n",
      "撒了32次星星\n",
      "31288\n",
      "撒了32次星星\n",
      "34089\n",
      "撒了32次星星\n",
      "36889\n",
      "撒了32次星星\n",
      "39689\n",
      "撒了32次星星\n",
      "42489\n",
      "撒了32次星星\n",
      "45289\n",
      "撒了32次星星\n",
      "48089\n",
      "撒了32次星星\n",
      "50889\n",
      "撒了32次星星\n",
      "53689\n",
      "撒了32次星星\n",
      "56489\n",
      "撒了32次星星\n",
      "59288\n",
      "撒了32次星星\n",
      "62088\n",
      "撒了32次星星\n",
      "64888\n",
      "撒了32次星星\n",
      "67689\n",
      "撒了32次星星\n",
      "70489\n",
      "撒了32次星星\n",
      "73289\n",
      "撒了32次星星\n",
      "76089\n",
      "撒了32次星星\n",
      "78889\n",
      "撒了32次星星\n",
      "81689\n",
      "撒了32次星星\n",
      "84489\n",
      "撒了32次星星\n",
      "87289\n",
      "撒了32次星星\n",
      "90089\n",
      "共计32条数据\n",
      "./test_data/Counting_Stars_EN_multi-evidence-retrieval-reasoning_128000_32_32.jsonl\n",
      "撒了32次星星\n",
      "22568\n",
      "撒了32次星星\n",
      "40925\n",
      "撒了32次星星\n",
      "59378\n",
      "撒了32次星星\n",
      "76151\n",
      "撒了32次星星\n",
      "93928\n",
      "撒了32次星星\n",
      "112510\n",
      "撒了32次星星\n",
      "129745\n",
      "撒了32次星星\n",
      "147424\n",
      "撒了32次星星\n",
      "164916\n",
      "撒了32次星星\n",
      "182440\n",
      "撒了32次星星\n",
      "199919\n",
      "撒了32次星星\n",
      "218068\n",
      "撒了32次星星\n",
      "235959\n",
      "撒了32次星星\n",
      "253914\n",
      "撒了32次星星\n",
      "271387\n",
      "撒了32次星星\n",
      "288485\n",
      "撒了32次星星\n",
      "306086\n",
      "撒了32次星星\n",
      "323889\n",
      "撒了32次星星\n",
      "341775\n",
      "撒了32次星星\n",
      "359259\n",
      "撒了32次星星\n",
      "377933\n",
      "撒了32次星星\n",
      "396289\n",
      "撒了32次星星\n",
      "414020\n",
      "撒了32次星星\n",
      "432135\n",
      "撒了32次星星\n",
      "449753\n",
      "撒了32次星星\n",
      "467551\n",
      "撒了32次星星\n",
      "485760\n",
      "撒了32次星星\n",
      "503200\n",
      "撒了32次星星\n",
      "520306\n",
      "撒了32次星星\n",
      "537684\n",
      "撒了32次星星\n",
      "555916\n",
      "撒了32次星星\n",
      "573853\n",
      "共计32条数据\n",
      "./test_data/Counting_Stars_ZH_multi-evidence-retrieval-reasoning_128000_32_32.jsonl\n",
      "撒了32次星星\n",
      "4159\n",
      "撒了32次星星\n",
      "6959\n",
      "撒了32次星星\n",
      "9759\n",
      "撒了32次星星\n",
      "12559\n",
      "撒了32次星星\n",
      "15359\n",
      "撒了32次星星\n",
      "18159\n",
      "撒了32次星星\n",
      "20959\n",
      "撒了32次星星\n",
      "23759\n",
      "撒了32次星星\n",
      "26559\n",
      "撒了32次星星\n",
      "29359\n",
      "撒了32次星星\n",
      "32158\n",
      "撒了32次星星\n",
      "34959\n",
      "撒了32次星星\n",
      "37759\n",
      "撒了32次星星\n",
      "40559\n",
      "撒了32次星星\n",
      "43359\n",
      "撒了32次星星\n",
      "46159\n",
      "撒了32次星星\n",
      "48959\n",
      "撒了32次星星\n",
      "51759\n",
      "撒了32次星星\n",
      "54559\n",
      "撒了32次星星\n",
      "57359\n",
      "撒了32次星星\n",
      "60158\n",
      "撒了32次星星\n",
      "62958\n",
      "撒了32次星星\n",
      "65758\n",
      "撒了32次星星\n",
      "68559\n",
      "撒了32次星星\n",
      "71359\n",
      "撒了32次星星\n",
      "74159\n",
      "撒了32次星星\n",
      "76959\n",
      "撒了32次星星\n",
      "79759\n",
      "撒了32次星星\n",
      "82559\n",
      "撒了32次星星\n",
      "85359\n",
      "撒了32次星星\n",
      "88159\n",
      "撒了32次星星\n",
      "90959\n",
      "共计32条数据\n"
     ]
    }
   ],
   "source": [
    "import uuid\n",
    "import pandas as pd\n",
    "import json\n",
    "import os\n",
    "import glob\n",
    "import jsonlines\n",
    "import requests\n",
    "from tqdm import trange\n",
    "import random\n",
    "import glob\n",
    "import re\n",
    "\n",
    "def get_context(language=\"EN\"):\n",
    "    if language == \"EN\":\n",
    "        context = \"\"\n",
    "        for file in glob.glob(\"./context_data/PaulGrahamEssays/*.txt\"):\n",
    "            with open(file, 'r') as f:\n",
    "                context += f.read().replace(\"\\n\", \" \")\n",
    "        return context\n",
    "    else:\n",
    "        string_punctuation = '!\"#$%&\\'()*+,-./:;<=>?@[\\\\]^_`{|}~'\n",
    "        context_file = open(\"./context_data/The_Story_of_the_Stone.txt\", \"r\", encoding=\"utf-8\")\n",
    "        context = \"\"\n",
    "        for i in context_file.readlines():\n",
    "            context += i.strip().replace(\"------------\",\" \").replace(\"\\n\", \" \").replace(\" \",\"\")\n",
    "        context = re.sub('[{}]'.format(string_punctuation),\"\",context)\n",
    "        context = re.sub('[a-zA-Z]','',context)\n",
    "        return context\n",
    "\n",
    "def get_stars(stars_dir, counting_times):\n",
    "    stars_file = open(stars_dir,\"r\")\n",
    "    return eval(stars_file.readline())[str(counting_times)]\n",
    "\n",
    "def sentence_with_star(language, test_type, indicator):\n",
    "    if language == \"ZH\":\n",
    "        if test_type == \"multi-evidence-retrieval-searching\":\n",
    "            single_star = f\"\\n小企鹅数了{a_stars[indicator]}颗★\\n\"\n",
    "        else:\n",
    "            single_star = f\"\\n小企鹅数了{r_stars[indicator]}颗★，但发现数错了，于是又数了一遍，这次数对了，是{a_stars[indicator]}颗★\\n\"\n",
    "    else:\n",
    "        if test_type == \"multi-evidence-retrieval-searching\":\n",
    "            single_star = f\"\\nThe little penguin counted {a_stars[indicator]} ★\\n\"\n",
    "        else:\n",
    "            single_star = f\"\\nThe little penguin counted {r_stars[indicator]} ★, but found that a mistake had been made, so the counting was done again, and this time {a_stars[indicator]} ★ was counted correctly.\\n\"\n",
    "    return single_star\n",
    "\n",
    "def select_question(language, test_type):\n",
    "    if language == \"ZH\":\n",
    "        searching_question = \"\\n\\n\\n\\n在这个月光皎洁、云雾缭绕的夜晚，小企鹅正望向天空，全神贯注地数★。请帮助小企鹅收集所数★的颗数，按照如下格式：{\\\"小企鹅\\\":[x,x,x,...]}，不要求和，[x,x,x,...]中数字为小企鹅每次数★的颗数，仅以JSON格式输出结果，不需要输出任何解释。\"\n",
    "        reasoning_question = \"\\n\\n\\n\\n在这个月光皎洁、云雾缭绕的夜晚，小企鹅正望向天空，全神贯注地数★。请帮助小企鹅收集所数★的正确颗数，按照如下格式：{\\\"小企鹅\\\":[x,x,x,...]}，不要求和，[x,x,x,...]中数字为小企鹅正确数★的颗数，仅以JSON格式输出结果，不需要输出任何解释。\"\n",
    "        if test_type == \"multi-evidence-retrieval-searching\":\n",
    "            return searching_question\n",
    "        else:\n",
    "            return reasoning_question\n",
    "    else:\n",
    "        searching_question = \"\\n\\n\\n\\n\" + \"On this moonlit and misty night, the little penguin is looking up at the sky and concentrating on counting ★. Please help the little penguin collect the number of ★, for example: {\\\"little_penguin\\\": [x, x, x,...]}. The summation is not required, and the numbers in [x, x, x,...] represent the counted number of ★ by the little penguin. Only output the results in JSON format without any explanation.\"\n",
    "        reasoning_question = \"\\n\\n\\n\\n\" + \"On this moonlit and misty night, the little penguin is looking up at the sky and concentrating on counting ★. Please help the little penguin collect the correct number of ★, for example: {\\\"little_penguin\\\": [x, x, x,...]}. The summation is not required, and the numbers in [x, x, x,...] represent the correctly counted number of ★ by the little penguin. Only output the results in JSON format without any explanation.\"\n",
    "        if test_type == \"multi-evidence-retrieval-searching\":\n",
    "            return searching_question\n",
    "        else:\n",
    "            return reasoning_question\n",
    "\n",
    "m = 32\n",
    "n = 32\n",
    "version = [[m, n]]\n",
    "language_types = [\"EN\", \"ZH\"]\n",
    "task_types = [\"multi-evidence-retrieval-searching\", \"multi-evidence-retrieval-reasoning\"]\n",
    "# MeRS-ZH, MeRS-EN, MeRR-EN, MeRR-ZH\n",
    "a_stars = get_stars(\"./context_data/a_stars.txt\", m)\n",
    "r_stars = get_stars(\"./context_data/r_stars.txt\", m)\n",
    "max_context_length = 128000\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    for task_type in task_types:\n",
    "        for language in language_types:\n",
    "            if language == \"EN\":\n",
    "                scalar = 0.8\n",
    "            else:\n",
    "                scalar = 0.7\n",
    "            context = get_context(language=language)\n",
    "            for m, n in version:\n",
    "                line_count = 0\n",
    "                interval = int(max_context_length/n)\n",
    "                context_size = [int(i*scalar) for i in range(interval, max_context_length+1, interval)]\n",
    "                file_name = f\"./test_data/Counting_Stars_{language}_{task_type}_{max_context_length}_{m}_{n}.jsonl\"\n",
    "                test_data = open(file_name, \"w\", encoding=\"utf-8\")\n",
    "                print(file_name)\n",
    "                for j in context_size:\n",
    "                    indicator = 0\n",
    "                    sprinkle_stars_context = \" \".join(context.split(\" \")[:j]) if language == \"EN\" else context[:j]\n",
    "                    for k in range(0, j, int(j / m)):\n",
    "                        single_star = sentence_with_star(language, task_type, indicator)\n",
    "                        if language == \"ZH\":\n",
    "                            sprinkle_stars_context = (sprinkle_stars_context[:k+int(j/m)+len(single_star)*indicator] + single_star + sprinkle_stars_context[k+int(j/m)+len(single_star)*indicator:])            \n",
    "                        else:\n",
    "                            sprinkle_stars_context = (\" \".join(sprinkle_stars_context.split(\" \")[:len(single_star.split(\" \"))*indicator+k+int(j / m)]) + single_star + \" \".join(sprinkle_stars_context.split(\" \")[int(j / m)+k+len(single_star.split(\" \"))*indicator:]))\n",
    "                        indicator += 1\n",
    "                        if indicator == m:\n",
    "                            print(f\"撒了{indicator}次星星\")\n",
    "                            break\n",
    "                    print(len(sprinkle_stars_context + select_question(language, task_type)))\n",
    "                    output_template = {\"question\": sprinkle_stars_context + select_question(language, task_type), \"context_size\": j, \"retrieval_question\": select_question(language, task_type),\n",
    "                                \"reference_counting_results\": a_stars, \"parameters\": {\"temperature\": 0.0}}\n",
    "                    print(json.dumps(output_template, ensure_ascii=False), file=test_data)\n",
    "                    line_count += 1\n",
    "                    test_data.flush()\n",
    "                test_data.close()\n",
    "                print(f\"共计{line_count}条数据\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}