82-NLP-生成任务GPT2数据预处理-项目实操-深度学习-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1c897996-60af-4524-b843-a8707894f455",
   "metadata": {},
   "source": [
    "### 基于GPT2的文本生成项目-数据预处理\n",
    "项目地址：https://github.com/yangjianxin1/GPT2-chitchat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a1414ef7-6aa2-475b-94ee-86d113e08123",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertTokenizerFast\n",
    "import argparse\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "import logging\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2ce94007-a849-42cf-b5b3-d9599e57c2e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_logger(log_path):\n",
    "    \"\"\"\n",
    "    将日志输出到日志文件和控制台\n",
    "    \"\"\"\n",
    "    logger = logging.getLogger(__name__)\n",
    "    logger.setLevel(logging.INFO)\n",
    "    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')\n",
    "    # 创建一个handler，用于写入日志文件\n",
    "    file_handler = logging.FileHandler(filename=log_path)\n",
    "    file_handler.setFormatter(formatter)\n",
    "    file_handler.setLevel(logging.INFO)\n",
    "    logger.addHandler(file_handler)\n",
    "\n",
    "    # 创建一个handler，用于将日志输出到控制台\n",
    "    console = logging.StreamHandler()\n",
    "    console.setLevel(logging.DEBUG)\n",
    "    console.setFormatter(formatter)\n",
    "    logger.addHandler(console)\n",
    "\n",
    "    return logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b6457910-16d9-4f28-8b98-db72e50a00ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess():\n",
    "    \"\"\"\n",
    "    对原始语料进行tokenize，将每段对话处理成如下形式：\"[CLS]utterance1[SEP]utterance2[SEP]utterance3[SEP]\"\n",
    "    \"\"\"\n",
    "    # 设置参数\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--vocab_path', default='E:/Fangwork/Dltools/GPT2-chitchat-master/vocab/vocab.txt', type=str, required=False,help='词表路径')\n",
    "    parser.add_argument('--log_path', default='E:/Fangwork/Dltools/GPT2-chitchat-master/data/preprocess.log', type=str, required=False, help='训练日志存放位置')\n",
    "    parser.add_argument('--train_path', default='E:/Fangwork/Dltools/GPT2-chitchat-master/data/train.txt', type=str, required=False, help='训练日志存放位置')\n",
    "    parser.add_argument('--save_path', default='E:/Fangwork/Dltools/GPT2-chitchat-master/data/train_51w.pkl', type=str, required=False, help='tokenize的训练数据集')\n",
    "    args =parser.parse_known_args()[0]\n",
    "\n",
    "    # 初始化日志对象\n",
    "    logger = create_logger(args.log_path)\n",
    "\n",
    "    # 初始化tokenizer\n",
    "    tokenizer = BertTokenizerFast(vocab_file=args.vocab_path, sep_token=\"[SEP]\", pad_token=\"[PAD]\", cls_token=\"[CLS]\")\n",
    "    sep_id = tokenizer.sep_token_id\n",
    "    cls_id = tokenizer.cls_token_id\n",
    "    logger.info(\"preprocessing data,data path:{}, save path:{}\".format(args.train_path, args.save_path))\n",
    "\n",
    "    # 读取训练数据集\n",
    "    with open(args.train_path, 'rb') as f:\n",
    "        data = f.read().decode(\"utf-8\")\n",
    "\n",
    "    # 需要区分linux和windows环境下的换行符\n",
    "    if \"\\r\\n\" in data:\n",
    "        train_data = data.split(\"\\r\\n\\r\\n\")\n",
    "    else:\n",
    "        train_data = data.split(\"\\n\\n\")\n",
    "    logger.info(\"there are {} dialogue in dataset\".format(len(train_data)))\n",
    "\n",
    "    # 开始进行tokenize\n",
    "    # 保存所有的对话数据,每条数据的格式为：\"[CLS]utterance1[SEP]utterance2[SEP]utterance3[SEP]\"\n",
    "    dialogue_len = []  # 记录所有对话tokenize之后的长度，用于统计中位数与均值\n",
    "    dialogue_list = []\n",
    "    with open(args.save_path, \"w\", encoding=\"utf-8\") as f:\n",
    "        for index, dialogue in enumerate(tqdm(train_data)):\n",
    "            if \"\\r\\n\" in data:\n",
    "                utterances = dialogue.split(\"\\r\\n\")\n",
    "            else:\n",
    "                utterances = dialogue.split(\"\\n\")\n",
    "\n",
    "            input_ids = [cls_id]  # 每个dialogue以[CLS]开头\n",
    "            for utterance in utterances:\n",
    "                input_ids += tokenizer.encode(utterance, add_special_tokens=False)\n",
    "                input_ids.append(sep_id)  # 每个utterance之后添加[SEP]，表示utterance结束\n",
    "            dialogue_len.append(len(input_ids))\n",
    "            dialogue_list.append(input_ids)\n",
    "    len_mean = np.mean(dialogue_len)\n",
    "    len_median = np.median(dialogue_len)\n",
    "    len_max = np.max(dialogue_len)\n",
    "    with open(args.save_path, \"wb\") as f:\n",
    "        pickle.dump(dialogue_list, f)\n",
    "    logger.info(\"finish preprocessing data,the result is stored in {}\".format(args.save_path))\n",
    "    logger.info(\"mean of dialogue len:{},median of dialogue len:{},max len:{}\".format(len_mean, len_median, len_max))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5bfc9f74-55ad-4751-8f07-da8a6221dae0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-05-12 16:09:55,454 - INFO - preprocessing data,data path:E:/Fangwork/Dltools/GPT2-chitchat-master/data/train.txt, save path:E:/Fangwork/Dltools/GPT2-chitchat-master/data/train_51w.pkl\n",
      "2022-05-12 16:09:55,471 - INFO - there are 9831 dialogue in dataset\n",
      "100%|████████████████████████████████████████████████████████████████████████████| 9831/9831 [00:07<00:00, 1354.13it/s]\n",
      "2022-05-12 16:10:02,794 - INFO - finish preprocessing data,the result is stored in E:/Fangwork/Dltools/GPT2-chitchat-master/data/train_51w.pkl\n",
      "2022-05-12 16:10:02,795 - INFO - mean of dialogue len:45.780998881090426,median of dialogue len:36.0,max len:1089\n"
     ]
    }
   ],
   "source": [
    "preprocess()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0030e263-91c8-4e4b-bb9c-4f84bcce2b21",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}