notebook/使用ChatHaruhi54K生成数据.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyN5GwB7Gdfblqr/Vmh/tFJ7",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Haruhi-2-Dev/blob/main/notebook/%E4%BD%BF%E7%94%A8ChatHaruhi54K%E7%94%9F%E6%88%90%E6%95%B0%E6%8D%AE.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "- [ ] RoleLLM的数据\n",
        "- [ ] 其他中文的hf的一些数据\n",
        "- [ ] ChatHaruhi54K本身的数据"
      ],
      "metadata": {
        "id": "GglE46SYH83z"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "这里我们假设输入数据已经组织成了hf格式\n",
        "\n",
        "可以有效载入到chatbot中"
      ],
      "metadata": {
        "id": "vPt6Zysa5H78"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip -q install openai tiktoken langchain chromadb datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "L6VrCMIb5T_j",
        "outputId": "db6895fa-c82b-426c-fdf5-5cea8add679a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m221.4/221.4 kB\u001b[0m \u001b[31m2.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m21.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m42.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m502.4/502.4 kB\u001b[0m \u001b[31m37.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m521.2/521.2 kB\u001b[0m \u001b[31m33.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.0/75.0 kB\u001b[0m \u001b[31m11.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m177.6/177.6 kB\u001b[0m \u001b[31m21.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m48.1/48.1 kB\u001b[0m \u001b[31m6.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m48.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m92.9/92.9 kB\u001b[0m \u001b[31m12.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.7/59.7 kB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.4/5.4 MB\u001b[0m \u001b[31m35.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.4/6.4 MB\u001b[0m \u001b[31m48.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.9/57.9 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m105.3/105.3 kB\u001b[0m \u001b[31m12.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.3/67.3 kB\u001b[0m \u001b[31m9.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m699.4/699.4 kB\u001b[0m \u001b[31m58.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m31.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m72.6/72.6 kB\u001b[0m \u001b[31m10.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m14.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m19.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m6.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.0/67.0 kB\u001b[0m \u001b[31m8.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.9/76.9 kB\u001b[0m \u001b[31m8.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m143.8/143.8 kB\u001b[0m \u001b[31m18.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m50.8/50.8 kB\u001b[0m \u001b[31m6.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m341.4/341.4 kB\u001b[0m \u001b[31m36.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.4/3.4 MB\u001b[0m \u001b[31m82.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m75.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m130.2/130.2 kB\u001b[0m \u001b[31m15.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m12.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for pypika (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "lida 0.0.10 requires kaleido, which is not installed.\n",
            "lida 0.0.10 requires python-multipart, which is not installed.\n",
            "llmx 0.0.15a0 requires cohere, which is not installed.\n",
            "tensorflow-probability 0.22.0 requires typing-extensions<4.6.0, but you have typing-extensions 4.8.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "组织数据"
      ],
      "metadata": {
        "id": "JNmjcFSeoNZF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"silk-road/ChatHaruhi-54K-Role-Playing-Dialogue\")\n",
        "\n",
        "role2tuple = {}\n",
        "\n",
        "from tqdm import tqdm\n",
        "for data in tqdm( dataset['train'] ):\n",
        "    agent_role = data['agent_role']\n",
        "    if agent_role not in role2tuple:\n",
        "        role2tuple[agent_role] = []\n",
        "\n",
        "    user_role = data['user_role']\n",
        "    user_question = data['user_question']\n",
        "\n",
        "    query = user_role + \":\" + user_question + \"\"\n",
        "\n",
        "    agent_response = data['agent_response']\n",
        "\n",
        "    target = agent_role + \":\" + agent_response + \"\"\n",
        "\n",
        "    # role2tuple[agent_role].append((query, target))\n",
        "\n",
        "    history_str = \"\"\n",
        "\n",
        "    tmp_data = {\n",
        "        \"query\": query,\n",
        "        \"target\": target,\n",
        "        \"history_str\": history_str,\n",
        "    }\n",
        "\n",
        "    role2tuple[agent_role].append(tmp_data)\n",
        "\n",
        "    more_dialogues = data['more_dialogues']\n",
        "\n",
        "    history_str += query + \"\\n\" + target\n",
        "\n",
        "    history_prefix = history_str\n",
        "\n",
        "    if len(more_dialogues) > 0:\n",
        "        n = len( more_dialogues )\n",
        "        for i in range(n):\n",
        "            if i == 0:\n",
        "                continue\n",
        "\n",
        "            sent = more_dialogues[i]\n",
        "            if len(sent) < len(agent_role):\n",
        "                continue\n",
        "\n",
        "            if sent.startswith(agent_role):\n",
        "                query = more_dialogues[i-1]\n",
        "                target = sent\n",
        "\n",
        "                history_str = history_prefix\n",
        "                for j in range(i-1):\n",
        "                    history_str += \"\\n\" + more_dialogues[j]\n",
        "\n",
        "                tmp_data = {\n",
        "                    \"query\": query,\n",
        "                    \"target\": target,\n",
        "                    \"history_str\": history_str,\n",
        "                }\n",
        "\n",
        "                role2tuple[agent_role].append(tmp_data)\n",
        "\n",
        "                # print(tuple_data)\n",
        "\n",
        "                # break\n",
        "\n",
        "        # break\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sh9u8M18oR8G",
        "outputId": "42f674a6-aab3-4944-f692-d710fc15c635"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 54726/54726 [00:04<00:00, 13667.94it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "tuple数据会保留在role2tuple"
      ],
      "metadata": {
        "id": "HTBjIkY4od7r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(role2tuple['Harry'][100])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-GVNjUkPofaQ",
        "outputId": "1a7c96e4-e49c-458b-a359-bdafca4fbbbb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'query': 'Voldemort:「You think you know more magic than I do? Than I, than Lord Voldemort, who has performed magic that Dumbledore himself never dreamed of?」', 'target': 'Harry:「Oh, he dreamed of it, but he knew more than you, knew enough not to do what you’ve done.」', 'history_str': ''}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "训练的参数设定"
      ],
      "metadata": {
        "id": "tl7H3ltbFaEJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "K_SEARCH = 5\n",
        "MAX_LEN_STORY = 1000 #这个是按照token算的\n",
        "MAX_LEN_HISTORY = 1200 # count with token"
      ],
      "metadata": {
        "id": "NFrw2aW_Fbsh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import openai\n",
        "key = \"key is not necessary here\"\n",
        "key_bytes = key.encode()\n",
        "os.environ[\"OPENAI_API_KEY\"] = key_bytes.decode('utf-8')"
      ],
      "metadata": {
        "id": "_fuSHAHU8-Kc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "%cd /content\n",
        "!rm -rf /content/Haruhi-2-Dev\n",
        "!git clone https://github.com/LC1332/Haruhi-2-Dev\n",
        "%cd /content/Haruhi-2-Dev"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "udVfoKmi9Tme",
        "outputId": "5f08a31d-3c69-48b4-9f55-247fc072c00b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content\n",
            "Cloning into 'Haruhi-2-Dev'...\n",
            "remote: Enumerating objects: 916, done.\u001b[K\n",
            "remote: Counting objects: 100% (223/223), done.\u001b[K\n",
            "remote: Compressing objects: 100% (116/116), done.\u001b[K\n",
            "remote: Total 916 (delta 150), reused 157 (delta 107), pack-reused 693\u001b[K\n",
            "Receiving objects: 100% (916/916), 105.96 MiB | 40.31 MiB/s, done.\n",
            "Resolving deltas: 100% (480/480), done.\n",
            "/content/Haruhi-2-Dev\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!ls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O6vNMO1f_E55",
        "outputId": "f7c89657-8dbd-4722-d481-6c0b368db889"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ChatHaruhi  data  LICENSE  notebook  Readme.md\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from ChatHaruhi import ChatHaruhi\n",
        "\n",
        "from ChatHaruhi.ChromaDB import ChromaDB\n",
        "\n",
        "class ChatHaruhiTrain(ChatHaruhi):\n",
        "\n",
        "    def build_story_db_from_vec( self, texts, vecs ):\n",
        "        self.db = ChromaDB()\n",
        "        self.stories = texts\n",
        "        self.db.init_from_docs( vecs, texts)\n",
        "\n",
        "    def add_story_with_expire(self, query, expire_story):\n",
        "        if self.db is None:\n",
        "            print(\"No vec DB！\")\n",
        "            return\n",
        "\n",
        "        query_vec = self.embedding(query)\n",
        "        stories = self.db.search(query_vec, self.k_search)\n",
        "\n",
        "        story_string = self.story_prefix_prompt + self.dialogue_divide_token\n",
        "\n",
        "        sum_story_token = self.tokenizer(story_string)\n",
        "\n",
        "        for story in stories:\n",
        "            if expire_story is not None and expire_story.strip() == story.strip():\n",
        "                continue\n",
        "\n",
        "            if expire_story == None and query.strip() in story.strip():\n",
        "                continue\n",
        "\n",
        "            story_token = self.tokenizer(story.strip()) + self.tokenizer(self.dialogue_divide_token)\n",
        "            if sum_story_token + story_token > self.max_len_story:\n",
        "                break\n",
        "            else:\n",
        "                sum_story_token += story_token\n",
        "                story_string += story.strip() + self.dialogue_divide_token\n",
        "\n",
        "        self.llm.user_message(story_string)\n",
        "\n",
        "    def generate_prompt(self, query, history_str, expire_story ):\n",
        "        # 这里修改下其他超参，不规范删了\n",
        "        # self.k_search = 5\n",
        "        # self.max_len_story = 1500\n",
        "        # self.max_len_history = 1200\n",
        "        self.story_prefix_prompt = \"\\nClassic scenes for the role are as follows:\"\n",
        "\n",
        "        self.llm.initialize_message()\n",
        "\n",
        "        self.llm.system_message(self.system_prompt)\n",
        "\n",
        "        self.add_story_with_expire(query, expire_story)\n",
        "\n",
        "        # self.add_history(history)\n",
        "        history_message = self.dialogue_divide_token + history_str.strip()\n",
        "        self.llm.user_message(history)\n",
        "\n",
        "        self.llm.user_message(query)\n",
        "\n",
        "        # self.llm.user_message(target)\n",
        "\n",
        "        return self.llm.messages\n",
        "\n",
        "    def add_history(self, history_list):\n",
        "\n",
        "        if len(history_list) == 0:\n",
        "            return\n",
        "\n",
        "        sum_history_token = 0\n",
        "        flag = 0\n",
        "        for history in history_list:\n",
        "            current_count = 0\n",
        "            if history is not None:\n",
        "                current_count += self.tokenizer(history)\n",
        "\n",
        "            sum_history_token += current_count\n",
        "            if sum_history_token > self.max_len_history:\n",
        "                break\n",
        "            else:\n",
        "                flag += 1\n",
        "\n",
        "        if flag == 0:\n",
        "            print('warning! no history added. the last dialogue is too long.')\n",
        "\n",
        "        # 是否添加历史前缀，\n",
        "        history_message = \"\"\n",
        "        for history in history_list[-flag:]:\n",
        "            history_message += history\n",
        "        self.llm.user_message(history_message)"
      ],
      "metadata": {
        "id": "WcWlOGEv9MGd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "这个时候我们一次性把所有的RoleLLM都先申请出来"
      ],
      "metadata": {
        "id": "2UC3bBgzopiN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "role_name_Haruhiu = {'汤师爷': 'tangshiye', 'tangshiye': 'tangshiye', 'Tangshiye': 'tangshiye',\n",
        "                     '慕容复': 'murongfu', 'murongfu': 'murongfu', 'Murongfu': 'murongfu',\n",
        "                     '李云龙': 'liyunlong', 'liyunlong': 'liyunlong', 'Liyunlong': 'liyunlong',\n",
        "                     'Luna': 'Luna', '王多鱼': 'wangduoyu', 'wangduoyu': 'wangduoyu',\n",
        "                     'Wangduoyu': 'wangduoyu', 'Ron': 'Ron', '鸠摩智': 'jiumozhi',\n",
        "                     'jiumozhi': 'jiumozhi', 'Jiumozhi': 'jiumozhi', 'Snape': 'Snape',\n",
        "                     '凉宫春日': 'haruhi', 'haruhi': 'haruhi', 'Haruhi': 'haruhi',\n",
        "                     'Malfoy': 'Malfoy', '虚竹': 'xuzhu', 'xuzhu': 'xuzhu',\n",
        "                     'Xuzhu': 'xuzhu', '萧峰': 'xiaofeng',\n",
        "                     'xiaofeng': 'xiaofeng', 'Xiaofeng': 'xiaofeng', '段誉': 'duanyu',\n",
        "                     'duanyu': 'duanyu', 'Duanyu': 'duanyu', 'Hermione': 'Hermione',\n",
        "                     'Dumbledore': 'Dumbledore', '王语嫣': 'wangyuyan', 'wangyuyan':\n",
        "                     'wangyuyan', 'Wangyuyan': 'wangyuyan', 'Harry': 'Harry',\n",
        "                     'McGonagall': 'McGonagall', '白展堂': 'baizhantang',\n",
        "                     'baizhantang': 'baizhantang', 'Baizhantang': 'baizhantang',\n",
        "                     '佟湘玉': 'tongxiangyu', 'tongxiangyu': 'tongxiangyu',\n",
        "                     'Tongxiangyu': 'tongxiangyu', '郭芙蓉': 'guofurong',\n",
        "                     'guofurong': 'guofurong', 'Guofurong': 'guofurong', '流浪者': 'wanderer',\n",
        "                     'wanderer': 'wanderer', 'Wanderer': 'wanderer', '钟离': 'zhongli',\n",
        "                     'zhongli': 'zhongli', 'Zhongli': 'zhongli', '胡桃': 'hutao', 'hutao': 'hutao',\n",
        "                     'Hutao': 'hutao', 'Sheldon': 'Sheldon', 'Raj': 'Raj',\n",
        "                     'Penny': 'Penny', '韦小宝': 'weixiaobao', 'weixiaobao': 'weixiaobao',\n",
        "                     'Weixiaobao': 'weixiaobao', '乔峰': 'qiaofeng', 'qiaofeng': 'qiaofeng',\n",
        "                     'Qiaofeng': 'qiaofeng', '神里绫华': 'ayaka', 'ayaka': 'ayaka',\n",
        "                     'Ayaka': 'ayaka', '雷电将军': 'raidenShogun', 'raidenShogun': 'raidenShogun',\n",
        "                     'RaidenShogun': 'raidenShogun', '于谦': 'yuqian', 'yuqian': 'yuqian',\n",
        "                     'Yuqian': 'yuqian', 'Professor McGonagall': 'McGonagall',\n",
        "                     'Professor Dumbledore': 'Dumbledore'}"
      ],
      "metadata": {
        "id": "vFc6Ol0vozyG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "count = 0\n",
        "\n",
        "for role in role2tuple:\n",
        "    if role in role_name_Haruhiu:\n",
        "        for ele in role2tuple[role]:\n",
        "            if isinstance( ele, dict):\n",
        "                count += 1\n",
        "\n",
        "print(count)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Bkko7Tg-o2SU",
        "outputId": "d614cd8d-00cb-41b8-c58c-43ec77af566b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "62663\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "新建各自的chatbot"
      ],
      "metadata": {
        "id": "sr0Hsqy5pXND"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import zipfile\n",
        "import os\n",
        "\n",
        "try:\n",
        "    os.makedirs(\"characters_zip\")\n",
        "except:\n",
        "    pass\n",
        "try:\n",
        "    os.makedirs(\"characters\")\n",
        "except:\n",
        "    pass\n",
        "\n",
        "role_en2bots = {}\n",
        "\n",
        "\n",
        "from tqdm import tqdm\n",
        "\n",
        "for ai_role_en in tqdm( role_name_Haruhiu.values() ):\n",
        "    if ai_role_en in role_en2bots:\n",
        "        continue\n",
        "\n",
        "    role_en2bots[ai_role_en] = ChatHaruhiTrain(role_name = ai_role_en)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n0NqGGTBpZWf",
        "outputId": "30ed79f0-413b-462d-c65b-911fa6726f33"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 76/76 [00:05<00:00, 14.27it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "对RoleLLM的数据进行组织"
      ],
      "metadata": {
        "id": "KzIajWcI7C-U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.schema import (\n",
        "    AIMessage,\n",
        "    HumanMessage,\n",
        "    SystemMessage\n",
        ")"
      ],
      "metadata": {
        "id": "abWj2YNQCFFc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(dir(chatbot.db))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5VroQwaBrm3x",
        "outputId": "ef9d18d2-4abb-487b-8530-a57e3d8711d9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['__abstractmethods__', '__annotations__', '__class__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__slots__', '__str__', '__subclasshook__', '__weakref__', '_abc_impl', 'client', 'collection', 'init_db', 'init_from_docs', 'load', 'path', 'save', 'search']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "'add', 'construct', 'copy', 'count', 'database', 'delete', 'dict', 'from_orm', 'get', 'id', 'json', 'metadata', 'modify', 'name', 'parse_file', 'parse_obj', 'parse_raw', 'peek', 'query', 'schema', 'schema_json', 'tenant', 'update', 'update_forward_refs', 'upsert', 'validate']\n",
        "[21]\n",
        "0s\n"
      ],
      "metadata": {
        "id": "93bo-Z5usWnc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(chatbot.db.collection.count())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3nnt2EGksKVE",
        "outputId": "83d29470-75e7-493d-8bf9-d82cfb9b5a07"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1048\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "result = chatbot.db.collection.get()\n",
        "print(result.keys())\n",
        "print(len(result['documents']))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "317fXU2Kso53",
        "outputId": "5b4a2b10-33a2-423b-fd51-9a18cf9a70df"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dict_keys(['ids', 'embeddings', 'metadatas', 'documents', 'uris', 'data'])\n",
            "1048\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5lynZ4m35BGy",
        "outputId": "13ec6f29-9d11-4d16-d8e2-ad0d2976cd1d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 3333/3333 [05:40<00:00,  9.80it/s]\n",
            "100%|██████████| 2834/2834 [04:56<00:00,  9.55it/s]\n",
            "100%|██████████| 2055/2055 [03:32<00:00,  9.67it/s]\n",
            "100%|██████████| 1317/1317 [02:09<00:00, 10.14it/s]\n",
            "100%|██████████| 624/624 [01:03<00:00,  9.78it/s]\n",
            "100%|██████████| 1191/1191 [01:59<00:00,  9.98it/s]\n",
            "100%|██████████| 1569/1569 [02:42<00:00,  9.64it/s]\n",
            "100%|██████████| 2569/2569 [04:25<00:00,  9.69it/s]\n",
            "100%|██████████| 1060/1060 [01:50<00:00,  9.61it/s]\n",
            "100%|██████████| 1850/1850 [03:09<00:00,  9.76it/s]\n",
            "100%|██████████| 1188/1188 [02:00<00:00,  9.87it/s]\n",
            "100%|██████████| 356/356 [00:35<00:00,  9.92it/s]\n",
            "100%|██████████| 3864/3864 [06:48<00:00,  9.46it/s]\n",
            "100%|██████████| 1255/1255 [02:05<00:00,  9.98it/s]\n",
            "100%|██████████| 3849/3849 [07:04<00:00,  9.08it/s]\n",
            "100%|██████████| 1400/1400 [02:23<00:00,  9.76it/s]\n",
            "100%|██████████| 746/746 [01:15<00:00,  9.82it/s]\n",
            "100%|██████████| 6226/6226 [10:43<00:00,  9.68it/s]\n",
            "100%|██████████| 1671/1671 [02:56<00:00,  9.48it/s]\n",
            "100%|██████████| 6311/6311 [11:00<00:00,  9.55it/s]\n",
            "100%|██████████| 1943/1943 [03:17<00:00,  9.82it/s]\n",
            "100%|██████████| 1188/1188 [02:03<00:00,  9.60it/s]\n",
            "100%|██████████| 652/652 [01:06<00:00,  9.83it/s]\n",
            "100%|██████████| 2185/2185 [03:47<00:00,  9.62it/s]\n",
            "100%|██████████| 939/939 [01:35<00:00,  9.81it/s]\n",
            "100%|██████████| 804/804 [01:27<00:00,  9.20it/s]\n",
            "100%|██████████| 1629/1629 [02:48<00:00,  9.67it/s]\n",
            "100%|██████████| 1725/1725 [02:54<00:00,  9.89it/s]\n",
            "100%|██████████| 1091/1091 [01:50<00:00,  9.86it/s]\n",
            "100%|██████████| 3423/3423 [05:47<00:00,  9.85it/s]\n",
            "100%|██████████| 1816/1816 [03:06<00:00,  9.73it/s]\n"
          ]
        }
      ],
      "source": [
        "from datasets import load_dataset\n",
        "from tqdm import tqdm\n",
        "\n",
        "save_datas = []\n",
        "\n",
        "# for role_name in tqdm(role_from_roleLLM):\n",
        "\n",
        "for role_name_zh in role2tuple.keys():\n",
        "    if role_name_zh not in role_name_Haruhiu:\n",
        "        continue\n",
        "    role_name_en = role_name_Haruhiu[role_name_zh]\n",
        "\n",
        "    chatbot = role_en2bots[role_name_en]\n",
        "\n",
        "\n",
        "    chatbot.k_search = K_SEARCH\n",
        "    chatbot.max_len_story = MAX_LEN_STORY\n",
        "    chatbot.max_len_history = MAX_LEN_HISTORY\n",
        "\n",
        "    all_tuples = role2tuple[role_name_zh]\n",
        "\n",
        "    # break\n",
        "\n",
        "    for tuple_data in tqdm(all_tuples):\n",
        "        query = tuple_data['query']\n",
        "        target = tuple_data['target']\n",
        "        history = tuple_data['history_str']\n",
        "        messages = chatbot.generate_prompt(query, history, None)\n",
        "        # print(prompt)\n",
        "\n",
        "        prompt = \"\"\n",
        "\n",
        "        for msg in messages:\n",
        "            if isinstance(msg, HumanMessage):\n",
        "                prompt += msg.content + \"\\n\"\n",
        "            elif isinstance(msg, AIMessage):\n",
        "                prompt += msg.content + \"\\n\"\n",
        "            elif isinstance(msg, SystemMessage):\n",
        "                prompt += msg.content + \"\\n\"\n",
        "\n",
        "        save_data = {\n",
        "            'context': prompt,\n",
        "            'target': target\n",
        "        }\n",
        "\n",
        "        save_datas.append(save_data)\n",
        "\n",
        "        # break\n",
        "\n",
        "    # break"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(all_tuples[2].keys())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "53ojkk0dtmAd",
        "outputId": "3354cb05-c9bf-4075-bfcf-34b2d8e854ec"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dict_keys(['query', 'target', 'history_str'])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(all_tuples[3])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cGWquxucmM49",
        "outputId": "8f75e189-3c60-48c5-e85f-06f243fe0c2b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'query': '陈近南:「大家是好朋友，这事虽然干系不小，却也不能相瞒。混在宫里当小太监的，就是我那小徒韦小宝自己。小宝，你出来见过众位前辈。」', 'target': '韦小宝:「哈哈哈，原来是众位前辈啊！小宝在宫里混得风生水起，怎么会忘记了这帮好兄弟呢？来，来，让我好好瞧瞧你们。」', 'history_str': ''}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(save_data['context'])\n",
        "\n",
        "print()\n",
        "\n",
        "print(save_data['target'])"
      ],
      "metadata": {
        "id": "nxYt8Of0L5kN",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "502c48a1-2141-49ed-a27b-77a8e474ef27"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "I want you to act like 韦小宝 from 鹿鼎记.\n",
            "If others‘ questions are related with the novel, please try to reuse the original lines from the novel.\n",
            "I want you to respond and answer like 韦小宝 using the tone, manner and vocabulary 韦小宝 would use. \n",
            "You must know all of the knowledge of 韦小宝.\n",
            "注意韦小宝是狡诈、机智、善于应变的\n",
            "对话中常常不拘礼节，不怕得罪他人，有时甚至带有一些粗鲁和挑衅\n",
            "常常以发财为目标，利用各种机会来获取财富。\n",
            "\n",
            "\n",
            "Classic scenes for the role are as follows:\n",
            "###\n",
            "李自成:「大丈夫恩怨分明，那日你师父没杀我，今日我也饶你一命。自今而后，你再向我女儿看上一眼、说一句话，我把你全身砸成了肉酱」\n",
            "韦小宝:「大丈夫一言既出，那就怎样。那日在三圣庵里，你和你的姘头陈圆圆，已将阿珂许配我为妻，难道又想赖么。你不许我向自己老婆看上一眼，说一句话，天下哪有这样的岳父大人」\n",
            "阿珂:「爹，咱们走，别理这小子胡说八道。他狗嘴里长不出象牙，有甚么好话说了」\n",
            "韦小宝:「好啊，你终于认了他啦。这父母之命，你听是不听」\n",
            "###\n",
            "阿珂:「甚么你说我娘有危险么」\n",
            "韦小宝:「你娘倒没危险，我却有大大的危险」\n",
            "阿珂:「怎么危险到你身上了」\n",
            "韦小宝:「胡大哥跟我八拜之交，是结义兄弟。倘若他在兵荒马乱之中，却跟你娘搂搂抱抱，勾勾搭搭，可不是做了我的岳父吗这辈份是一塌胡涂了」\n",
            "###\n",
            "阿珂:「不不要郑郑公子是你么」\n",
            "旁白:韦小宝想\"你做梦也梦到郑公子，只道是他爬上了你床，好快活么\"\n",
            "韦小宝:「是我」\n",
            "阿珂:「不，不你不要」\n",
            "郑克爽:「阿珂，阿珂，你在哪里」\n",
            "阿珂:「你你是谁怎么我我」\n",
            "韦小宝:「是你的亲老公，你也听不出」\n",
            "阿珂:「我在这里放开手小鬼，你干干甚么」\n",
            "郑克爽:「甚么」\n",
            "韦小宝:「好师弟，求求你，快放开我」\n",
            "韦小宝:「我在床上，抱着我老婆。我在洞房花烛，你来干甚么要闹新房么」\n",
            "###\n",
            "阿珂:「你你怎么不早说他又说什么」\n",
            "韦小宝:「他说，这几位侠女要到台湾去玩玩，他就带她们同去，说要尽什么地主之之什么的。」\n",
            "阿珂:「地主之谊。」\n",
            "韦小宝:「对了，对了原来师姊刚才跟在我后面，都听见了。」\n",
            "阿珂:「可惜什么」\n",
            "韦小宝:「可惜不是郑公子追上来」\n",
            "阿珂:「哇」\n",
            "九难:「小宝，别老是使坏，激你师姊」\n",
            "韦小宝:「是，是」\n",
            "韦小宝:「是，是」\n",
            "###\n",
            "\n",
            "\n",
            "阿珂:「是，是我求你的。你是英雄好汉，大丈夫挺身而出，济人之急，又又最听我话的。」\n",
            "\n",
            "\n",
            "韦小宝:「师姊，我对你一番心意，你现在总明白了。不论你叫我做什么事，我都一口答应，不会皱一皱眉头。你既要我拜堂成亲，我自然答应。」\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "NwG1ZHv7MZ9H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "save_datas是一个list of dict\n",
        "save_name = \"/content/ChatHaruhi_rolellm_fusion.jsonl\"\n",
        "\n",
        "帮我用utf-8编码，jsonl格式，保存save_datas中的数据\n",
        "ensure_ascii = False"
      ],
      "metadata": {
        "id": "tgQ19MI9FEkn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(len(save_datas))"
      ],
      "metadata": {
        "id": "OkRms1zVF4bs",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2a49ddd7-e920-49ed-bb59-01a1d8dc2157"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "13166\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "save_name = \"/content/ChatHaruhi_54K_retide.jsonl\"\n",
        "\n",
        "with open(save_name, 'w', encoding='utf8') as f:\n",
        "    for data in save_datas:\n",
        "        json.dump(data, f, ensure_ascii=False)\n",
        "        f.write('\\n')"
      ],
      "metadata": {
        "id": "3wW5RXl55VK3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "U_nKXIs957Oe"
      }
    }
  ]
}