preprocess.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GBK->UTF8 (Don't run)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "\n",
    "# rootdir = 'data/'\n",
    "# if os.path.exists('utf-8'):\n",
    "#   os.system('rm -rf utf-8')\n",
    "# for subdir, dirs, files in os.walk(rootdir):\n",
    "#   for i,file in enumerate(files):\n",
    "#       orignal_filename = os.path.join(subdir, file)\n",
    "#       new_filename = os.path.join('\\'utf-8',orignal_filename)+'\\''\n",
    "#       #new_filename = os.path.join('utf-8',subdir, str(i)+\".txt\")\n",
    "#       os.system(\"mkdir -p `dirname \"+new_filename+\"`\")\n",
    "#       os.system(\"iconv -c -f GB2312 -t UTF-8 \\'\"+orignal_filename+'\\' > '+new_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Annoated->Json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'utf-8/data/people/annotated/6-1/6_052\\xe4\\xbc\\x9a\\xe8\\xaf\\xb4\\xe8\\xaf\\x9d\\xe7\\x9a\\x84\\xe7\\x9c\\xbc\\xe7\\x9d\\x9b.txt'\n",
      "b'utf-8/data/people/annotated/8-1/8_205\\xe6\\x88\\x91\\xe4\\xb8\\x8e\\xe7\\x88\\xb6\\xe4\\xba\\xb2.txt'\n",
      "b'utf-8/data/people/annotated/9-1/9_105\\xe5\\xa4\\x96\\xe5\\xa9\\x86\\xe7\\x9a\\x84\\xe8\\xae\\xb0\\xe5\\xbf\\x86-.txt'\n",
      "b'utf-8/data/people/annotated/9-1/9_104\\xe5\\xa4\\x96\\xe5\\xa9\\x86\\xe7\\x9a\\x84\\xe6\\x89\\x8b.txt'\n",
      "b'utf-8/data/people/annotated/9-1/9_553\\xe8\\xae\\xa9\\xe6\\x88\\x91\\xe5\\xbf\\x83\\xe5\\xad\\x98\\xe6\\x84\\x9f\\xe6\\xbf\\x80\\xe7\\x9a\\x84\\xe4\\xba\\xba3.txt'\n",
      "b'utf-8/data/people/annotated/9-1/9_046\\xe4\\xbd\\xa0\\xe6\\x98\\xaf\\xe6\\x88\\x91\\xe5\\xbf\\x83\\xe4\\xb8\\xad\\xe7\\x9a\\x84\\xe6\\xad\\x8c.txt'\n",
      "b'utf-8/data/people/annotated/9-1/9_408\\xe6\\x9d\\x91\\xe9\\x95\\xbf\\xe5\\xb0\\xbd\\xe5\\x9b\\xb9\\xe5\\x9c\\x84.txt'\n",
      "b'utf-8/data/people/annotated/9-1/9_560\\xe8\\xae\\xa9\\xe6\\x88\\x91\\xe8\\xbd\\xbb\\xe8\\xbd\\xbb\\xe5\\x9c\\xb0\\xe5\\x91\\x8a\\xe8\\xaf\\x89\\xe4\\xbd\\xa0.txt'\n",
      "b'utf-8/data/people/annotated/9-1/9_400\\xe6\\x9c\\x8b\\xe5\\x8f\\x8b (2).txt'\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "import json\n",
    "import re,os\n",
    "rootdir = b'utf-8/data/people/annotated'\n",
    "parse_1 = re.compile(r'\\$<*(.+)\\[')\n",
    "parse_1_1 = re.compile(r'\\$<*(.+)=')\n",
    "parse_1_2 = re.compile(r'\\$<*(.+)>\\$')\n",
    "parse_2 = re.compile(r'\\[(.+)\\]')\n",
    "parse_3 = re.compile(r'=(.*)>')\n",
    "parse_4 = re.compile(r'<(.*?)>')\n",
    "parse_5 = re.compile(r'\\$<(.*?)>\\$')\n",
    "parse_6 = re.compile(r'#<(.*?)>#@')\n",
    "passages = []\n",
    "def removeEmpty(list_in):\n",
    "  list_out = []\n",
    "  for i in list_in:\n",
    "    if i.strip()!='':\n",
    "      list_out.append(i.strip())\n",
    "  if list_out == []:\n",
    "    list_out.append(\"NULL\")\n",
    "  return list_out\n",
    "\n",
    "def removeNum(list_in):\n",
    "  list_out = []\n",
    "  for i in list_in:\n",
    "    seg = i.split(\"<\")\n",
    "    seg = [re.sub(r'[0-9#$<>\\s]',\"\",x) for x in seg]\n",
    "    list_out.extend(seg)\n",
    "  return list_out\n",
    "  if list_out == []:\n",
    "    list_out.append(\"NULL\")\n",
    "  return list_out\n",
    "for subdir, dirs, files in os.walk(rootdir):\n",
    "  for ttt,file in enumerate(files):\n",
    "      paragraphs = []\n",
    "      paragraph = {}\n",
    "      paragraph_groups = []\n",
    "      sentences = []\n",
    "      passage = {}\n",
    "      \n",
    "      filename = os.path.join(subdir, file)\n",
    "      f = open(filename,'r',encoding='utf-8')\n",
    "      first_line = f.readline().strip()\n",
    "      while first_line == \"\":\n",
    "        first_line = f.readline().strip()\n",
    "      try:\n",
    "        if '[' not in first_line:\n",
    "          gender = \"NULL\"\n",
    "        else:\n",
    "          gender = re.findall(parse_2,first_line)[0]\n",
    "          \n",
    "        if '=' not in first_line:\n",
    "          qualities = [\"NULL\"]\n",
    "        else:\n",
    "          qualities = re.findall(parse_3,first_line)[0].split(\"+\")\n",
    "        if gender == \"NULL\":\n",
    "          if qualities ==[\"NULL\"]:\n",
    "            figure = re.findall(parse_1_2,first_line)[0]\n",
    "          else:\n",
    "            figure = re.findall(parse_1_1,first_line)[0]\n",
    "        else:\n",
    "            figure = re.findall(parse_1,first_line)[0]\n",
    "        second_line = f.readline().strip()\n",
    "        while second_line == \"\":\n",
    "          second_line = f.readline().strip()\n",
    "        if \"#\" in second_line:\n",
    "          sentiments = re.findall(parse_4,second_line)\n",
    "          third_line = f.readline().strip()\n",
    "          while third_line == \"\":\n",
    "            third_line = f.readline().strip()\n",
    "        else:\n",
    "          third_line = second_line\n",
    "          sentiments = [\"NULL\"]\n",
    "        title = third_line.strip()\n",
    "        fourth_line = f.readline().strip()\n",
    "        while fourth_line == \"\":\n",
    "            fourth_line = f.readline().strip()\n",
    "        if \"总\" in fourth_line and \"分\" in fourth_line:\n",
    "          structure = re.findall(parse_4,fourth_line)[0]\n",
    "        else:\n",
    "          structure = \"NULL\"\n",
    "      except:\n",
    "        print(filename)\n",
    "        continue\n",
    "      qualities = removeEmpty(qualities)\n",
    "      sentiments = removeEmpty(sentiments)\n",
    "      passage['figure'] = figure.strip() #if \"陌生人[\" not in p['figure'] else \"陌生人\"\n",
    "      passage['gender'] = gender.strip()   #if '[' not in p['gender'] else p['gender'][-1]\n",
    "      passage['qualities'] = qualities\n",
    "      passage['sentiments'] = sentiments\n",
    "      passage['title'] = title.strip()\n",
    "      passage['structure'] = structure.strip()\n",
    "      \n",
    "      if structure == \"NULL\":\n",
    "        passage_first_line = fourth_line\n",
    "      else:\n",
    "        passage_first_line = f.readline()\n",
    "      para_info = re.findall(parse_5,passage_first_line)\n",
    "      para_info = removeEmpty(para_info)\n",
    "      para_info = removeNum(para_info)\n",
    "      paragraph_group = {\"paraGroupFunc\":para_info}\n",
    "      para_index = 0\n",
    "      for index,line in enumerate(f):\n",
    "        para_info = re.findall(parse_5,line)\n",
    "        if para_info:\n",
    "          if paragraphs:\n",
    "            paragraph_group['paraGroup'] = paragraphs\n",
    "            paragraph_groups.append(copy.deepcopy(paragraph_group))\n",
    "          paragraphs = []\n",
    "          para_info = removeEmpty(para_info)\n",
    "          para_info = removeNum(para_info)\n",
    "          paragraph_group = {\"paraGroupFunc\":para_info}\n",
    "        else:\n",
    "          sentences = []\n",
    "          parse_add_and = re.compile(r\"(#)[\\u4e00-\\u9fa5]\")\n",
    "          line = re.sub(parse_add_and,\"#@\",line)\n",
    "          splits = re.split(parse_6,line.strip())\n",
    "          annotations = re.findall(parse_6,line)\n",
    "          sen_annotate = []\n",
    "          for i in range(len(splits)):\n",
    "            if splits[i] == \"\":\n",
    "              continue\n",
    "            if splits[i] not in annotations:\n",
    "              splits[i] = splits[i].replace(\"。。。。。。\",\"......\")\n",
    "              splits[i] = splits[i].replace(\"……\",\"......\")\n",
    "              splits[i] = splits[i].replace(\"``````\",\"......\")\n",
    "              splits[i] = splits[i].replace(\"”\",\"\\\"\")\n",
    "              splits[i] = splits[i].replace(\"“\",\"\\\"\")\n",
    "              splits[i] = splits[i].replace(\"＂\",\"\\\"\")\n",
    "              splits[i] = splits[i].replace(\"DD\",\"--\")\n",
    "              sentence = {\"senGroupFunc\":sen_annotate,\"senGroup\":splits[i].strip()}\n",
    "              sentences.append(copy.deepcopy(sentence))\n",
    "              sen_annotate = []\n",
    "            else:\n",
    "              sen_annotate.append(splits[i])\n",
    "          if sentences == []:\n",
    "            continue\n",
    "          paragraph[\"paraIndex\"] = para_index\n",
    "          paragraph[\"paragraph\"] = sentences\n",
    "          para_index+=1\n",
    "          paragraphs.append(copy.deepcopy(paragraph))\n",
    "      paragraph_group['paraGroup'] = paragraphs\n",
    "      paragraph_groups.append(copy.deepcopy(paragraph_group))\n",
    "      passage['filename'] = filename.decode('utf-8')\n",
    "      passage['paraGroups'] = paragraph_groups\n",
    "      passages.append(passage)\n",
    "\n",
    "def sgf_split(sgf_list,sen_list,current_index):\n",
    "  output_list = []\n",
    "  num_pattern = re.compile(r'([1-9])')\n",
    "  for i in sgf_list:\n",
    "    seg = i.split(\"<\")\n",
    "    seg = [re.sub(r'[#$<>\\s]',\"\",x) for x in seg]\n",
    "    for s in seg:\n",
    "      para = re.sub('([^1-9])([1-9])', r\"\\1\\n\\2\", s).split('\\n')\n",
    "      if len(para)==1:\n",
    "        para.append(len(sen_list))\n",
    "\n",
    "      para = (para[0],list(range(current_index,int(para[1])+current_index)))\n",
    "      output_list.append(para)\n",
    "      if len(para[1])<len(sen_list):\n",
    "        after_para = (para[0]+\"_after\", list(range(para[1][-1]+1,current_index+len(sen_list))))\n",
    "        output_list.append(after_para)\n",
    "  return output_list\n",
    "def cut_sent(para):\n",
    "    para = re.sub('\\?\\?\\?\\?\\?\\?',\"\\?\",para)\n",
    "    para = re.sub('\\(\\s\\)','',para)\n",
    "    para = re.sub('([。！!?\\?])([^\\\"”＂’\\'。！!?\\?])', r\"\\1\\n\\2\", para)  # 单字符断句符\n",
    "    para = re.sub('(\\.{6})([^\\\"”＂’\\'])', r\"\\1\\n\\2\", para)  # 英文省略号\n",
    "    para = re.sub('(\\…{2})([^\\\"”＂’\\'])', r\"\\1\\n\\2\", para)  # 中文省略号\n",
    "    para = re.sub('([。！!？\\?][\\\"”＂’\\'])([^，。！？\\?])', r'\\1\\n\\2', para)\n",
    "    # 如果双引号前有终止符，那么双引号才是句子的终点，把分句符\\n放到双引号后，注意前面的几句都小心保留了双引号\n",
    "    para = para.rstrip()  # 段尾如果有多余的\\n就去掉它\n",
    "    # 很多规则中会考虑分号;，但是这里我把它忽略不计，破折号、英文双引号等同样忽略，需要的再做些简单调整即可。\n",
    "    return para.split(\"\\n\")\n",
    "\n",
    "for ttt,p in enumerate(passages):\n",
    "  sentences = []\n",
    "  annotations = []\n",
    "  paragraphs = []\n",
    "  paragraph_groups = []\n",
    "  for pg in p['paraGroups']:\n",
    "    previous_paragraph_num = len(paragraphs)\n",
    "    for para in pg['paraGroup']:\n",
    "      previous_sentences_num = len(sentences)\n",
    "      for sg in para['paragraph']:\n",
    "        annotation = sg['senGroupFunc']\n",
    "        sequence = sg['senGroup']\n",
    "        sen_list = cut_sent(sequence)\n",
    "        annotation_tuples = sgf_split(annotation,sen_list,len(sentences))\n",
    "        annotations.extend(annotation_tuples)\n",
    "        sentences.extend(sen_list)\n",
    "      paragraphs.append((para['paraIndex'],list(range(previous_sentences_num,len(sentences)))))\n",
    "    paragraph_groups.append((pg['paraGroupFunc'],list(range(previous_paragraph_num,len(paragraphs)))))\n",
    "  # quota\n",
    "  total_pair_quota = 0\n",
    "  for i in range(len(sentences)-1):\n",
    "    quota_count = len(re.findall(\"\\\"\",sentences[i]))\n",
    "    if quota_count==0:\n",
    "      continue\n",
    "    if quota_count%2!=0 and sentences[i][-1] == \"\\\"\" and total_pair_quota%2==0:\n",
    "      sentences[i+1]=\"\\\"\"+sentences[i+1]\n",
    "      sentences[i]=sentences[i][:-1]\n",
    "      total_pair_quota+=quota_count-1\n",
    "    else:\n",
    "      total_pair_quota+=quota_count\n",
    "  p['all_sentences'] = sentences\n",
    "  p['pair_annotation_sentences'] = annotations \n",
    "  p['pair_paragraph_sentences'] = paragraphs\n",
    "  p['pair_annotations_paragraphs'] = paragraph_groups\n",
    "with open('people.json', 'w',encoding='utf-8') as fout:\n",
    "  json.dump(passages,fout,indent=4,ensure_ascii=False)\n",
    "     #fout.write(json.dumps(passages, indent=4,ensure_ascii=False).encode('utf8'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train/Val/Test Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "passages = json.load(open('people.json', 'r',encoding='utf-8'))\n",
    "total_num = len(passages)\n",
    "np.random.seed(42)\n",
    "\n",
    "all_index = list(range(0,total_num))\n",
    "np.random.shuffle(all_index)\n",
    "train_index = all_index[:int(total_num*0.7)]\n",
    "valid_index = all_index[int(total_num*0.7):int(total_num*0.8)]\n",
    "test_index = all_index[int(total_num*0.8):]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_text_file = open(\"processed_data/train.txt\",'w',encoding='utf-8')\n",
    "valid_text_file = open(\"processed_data/valid.txt\",'w',encoding='utf-8')\n",
    "test_text_file = open(\"processed_data/test.txt\",\"w\",encoding='utf-8')\n",
    "for index,passage in enumerate(passages):\n",
    "  if index in train_index:\n",
    "    text_file = train_text_file \n",
    "  elif index in valid_index:\n",
    "    text_file = valid_text_file\n",
    "  else:\n",
    "    text_file = test_text_file\n",
    "  for pair_paragraph_sentences in passage['pair_paragraph_sentences']:\n",
    "    for sentence_index in pair_paragraph_sentences[1]:\n",
    "      text_file.write(passage['all_sentences'][sentence_index])\n",
    "      text_file.write('\\n')\n",
    "    text_file.write('\\n')\n",
    "  text_file.write('\\n')\n",
    "train_text_file.close()\n",
    "test_text_file.close()\n",
    "valid_text_file.close()\n",
    "\n",
    "random_poses = []\n",
    "test_text_file = open(\"/data/baihe/dataset/Chinese_Eassy/processed_data/test.txt\",\"r\",encoding='utf-8')\n",
    "test_input = open(\"/data/baihe/dataset/Chinese_Eassy/processed_data/test_partA.txt\",\"w\",encoding='utf-8')\n",
    "test_output = open(\"/data/baihe/dataset/Chinese_Eassy/processed_data/test_partB.txt\",\"w\",encoding='utf-8')\n",
    "text = test_text_file.read()\n",
    "passages =text.strip().split('\\n\\n\\n')\n",
    "for index,passage in enumerate(passages):\n",
    "  random_pos = np.random.randint(int(len(passage)*0.5))\n",
    "  while len(passage[:random_pos].split('\\n')) ==1 or passage[random_pos]=='\\n' or passage[random_pos-1]!='\\n':\n",
    "    random_pos+=1\n",
    "  random_poses.append(random_pos)\n",
    "#   while passage[random_pos]=='\\n':\n",
    "#     random_pos-=1\n",
    "  test_input.write(passage[:random_pos])\n",
    "  test_input.write(\"\\n\"+\"*\"*10+str(index)+\"*\"*10+'\\n')\n",
    "  test_output.write(passage[random_pos:])\n",
    "  test_output.write(\"\\n\"+\"*\"*10+str(index)+\"*\"*10+'\\n')\n",
    "test_input.close()\n",
    "test_output.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# State Sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# State of Sentence\n",
    "# State of Passage\n",
    "# State of Paragraph\n",
    "state = \"paragraph\"\n",
    "start_para = [\"总领全文\",\"概括介绍\",\"俗语开篇\",\"诗歌开篇\",\"歌词开篇\",\"故事开篇\",\"对话开篇\",\"题记开篇\",\"开门见山\", \"设置悬念\",\"背景介绍\"]\n",
    "mid_para = [\"过渡\" ,\"转折\",\"中心段\",\"点题照应\", \"人物描写\" , \"事件描写\" ]\n",
    "end_para = [\"首尾照应\",\"文题照应\",\"前后照应\",\"抒情\",\"主题升华\"]\n",
    "para_dict = {x:\"start\" for x in start_para}\n",
    "para_dict.update({x:\"end\" for x in end_para})\n",
    "para_dict.update({x:\"mid\" for x in mid_para})\n",
    "\n",
    "passages = json.load(open('people.json', 'r',encoding='utf-8'))\n",
    "train_state_file = open(\"processed_data/state_para/train.txt\",'w',encoding='utf-8')\n",
    "valid_state_file = open(\"processed_data/state_para/valid.txt\",'w',encoding='utf-8')\n",
    "test_state_file = open(\"processed_data/state_para/test.txt\",\"w\",encoding='utf-8')\n",
    "for index,passage in enumerate(passages):\n",
    "  if index in train_index:\n",
    "    state_file = train_state_file \n",
    "  elif index in valid_index:\n",
    "    state_file = valid_state_file\n",
    "  else:\n",
    "    state_file = test_state_file\n",
    "  \n",
    "  for pair_annotations_paragraphs in passage['pair_annotations_paragraphs']:\n",
    "    para_index_list = pair_annotations_paragraphs[1]\n",
    "    para_func_list = pair_annotations_paragraphs[0]\n",
    "    if para_func_list[0] not in para_dict.keys():\n",
    "      continue\n",
    "    for para_index in para_index_list: \n",
    "      sent_index_list = passage['pair_paragraph_sentences'][para_index][1]\n",
    "      for sent_index in sent_index_list:\n",
    "        state_file.write(passage['all_sentences'][sent_index].strip())\n",
    "      state_file.write(\"|||\")\n",
    "      state_file.write(para_dict[para_func_list[0]].strip())\n",
    "      #state_file.write(\" \".join(para_func_list))\n",
    "      state_file.write('\\n')\n",
    "  state_file.write('\\n')\n",
    "train_state_file.close()\n",
    "test_state_file.close()\n",
    "valid_state_file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = \"figure\"\n",
    "figure_count = dict()\n",
    "passages = json.load(open('people.json', 'r',encoding='utf-8'))\n",
    "train_state_file = open(\"processed_data/figure/train.txt\",'w',encoding='utf-8')\n",
    "valid_state_file = open(\"processed_data/figure/valid.txt\",'w',encoding='utf-8')\n",
    "test_state_file = open(\"processed_data/figure/test.txt\",\"w\",encoding='utf-8')\n",
    "for index,passage in enumerate(passages):\n",
    "  if \"[\" in passage['figure']:\n",
    "    passage['figure'] = passage['figure'][4:-1]\n",
    "  if passage['figure'] in figure_count.keys():\n",
    "    figure_count[passage['figure']]+=1\n",
    "  else:\n",
    "    figure_count[passage['figure']]=1\n",
    "  if index in train_index:\n",
    "    text_file = train_state_file \n",
    "  elif index in valid_index:\n",
    "    text_file = valid_state_file\n",
    "  else:\n",
    "    text_file = test_state_file\n",
    "  for pair_paragraph_sentences in passage['pair_paragraph_sentences']:\n",
    "    for sentence_index in pair_paragraph_sentences[1]:\n",
    "      text_file.write(passage['all_sentences'][sentence_index])\n",
    "      text_file.write(\" ||| \"+passage['figure'])\n",
    "      text_file.write('\\n')\n",
    "    text_file.write('\\n')\n",
    "  text_file.write('\\n')\n",
    "  \n",
    "train_state_file.close()\n",
    "test_state_file.close()\n",
    "valid_state_file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = \"prompt\"\n",
    "passages = json.load(open('people.json', 'r',encoding='utf-8'))\n",
    "train_state_file = open(\"processed_data/prompt/train.txt\",'w',encoding='utf-8')\n",
    "valid_state_file = open(\"processed_data/prompt/valid.txt\",'w',encoding='utf-8')\n",
    "test_state_file = open(\"processed_data/prompt/test.txt\",\"w\",encoding='utf-8')\n",
    "for index,passage in enumerate(passages):\n",
    "  if index in train_index:\n",
    "    text_file = train_state_file \n",
    "  elif index in valid_index:\n",
    "    text_file = valid_state_file\n",
    "  else:\n",
    "    text_file = test_state_file\n",
    "  for pair_paragraph_sentences in passage['pair_paragraph_sentences']:\n",
    "    for sentence_index in pair_paragraph_sentences[1]:\n",
    "      text_file.write(passage['all_sentences'][sentence_index])\n",
    "      text_file.write(\" ||| \"+passage['figure'])\n",
    "      text_file.write('\\n')\n",
    "    text_file.write('\\n')\n",
    "  text_file.write('\\n')\n",
    "  \n",
    "train_state_file.close()\n",
    "test_state_file.close()\n",
    "valid_state_file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Topic focus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = open('processed_data/figure/test.txt','r',encoding='utf-8').read().split('\\n\\n\\n')[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "topics = ['母亲', '同桌', '姥爷', '老师', '老妈', '同学', '父亲',\n",
    "          '姥姥', '哥哥', '清洁工', '弟弟', '爷爷', '陌生人', '奶奶',\n",
    "          '姐姐', '阿姨', '偶然遇到的人', '环卫工人', '班主任', '外婆', '邻居',\n",
    "          '妹妹', '乞丐', '表妹', '小偷', '伯父', '表弟', '爸爸', '表哥',\n",
    "          '修鞋师傅', '叔叔', '妈妈', '老爸', '舅舅', '姑姑', '捡破烂奶奶',\n",
    "          '外公',  '教练', '盲人', '老奶奶', '卖饼阿姨',\n",
    "          '售票员', '教官', '乘客', '补鞋老人', '墓地少女', '老妪', '擦皮鞋小女孩',\n",
    "          '交警', '邻居爷爷', '表姐', '婆婆', '川剧大师', '堂姐', '卖别针的老人', '少年',\n",
    "          '邮递员', '老太太', '捡破烂老人', '磨剪刀爷爷', '修车师傅', '小女孩', '祖母',\n",
    "          '老爷爷', '小摊主', '警察叔叔', '高师傅', '男老师', '太姥爷', '司机', '路人',\n",
    "          '阿婆', '狱警', '村长', '售货员', '祖父', '乘务员', '拾荒者', '电梯师傅',\n",
    "          '民工', '护士', '姨妈', '继父', '残疾人', '堂弟', '园丁', '运货员',\n",
    "          '学姐', '前男友', '消防队员', '外祖母', '侄子', '闺蜜',\n",
    "          '对手', '知己', '拾荒奶奶', '养父', '眼镜男', '姐妹', '太姥爷和太姥姥', '磨刀人',\n",
    "          '卖馍老头', '老夫妇', '干爸', '农村小姑娘', '太姥姥', '师姐', '食堂大妈',\n",
    "          '交通协管员', '小侄儿', '修车老爷爷', '太祖母', '买菜女', '校长', '废品老人']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_out = open('processed_data/figure/test_topic_label.txt','w',encoding='utf-8')\n",
    "f_all = open('processed_data/figure/all_topic.txt','w',encoding='utf-8')\n",
    "for t in topics:\n",
    "  f_all.write(t+'\\n')\n",
    "for line in f:\n",
    "  annotation = line.split('\\n')[0].split('|||')[1].strip()\n",
    "  f_out.write(annotation+'\\n')\n",
    "f_out.close()\n",
    "f_all.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re \n",
    "pattern = r\"\\*\" * 10 + \"[0-9]+\" + \"\\*\" * 10\n",
    "generated_text = open('processed_data/test_partB.txt','r',encoding='utf-8').read()\n",
    "labels = open('processed_data/figure/test_topic_label.txt','r',encoding='utf-8').read().strip().split('\\n') \n",
    "all_label = open('processed_data/figure/all_topic.txt','r',encoding='utf-8').read().strip().split('\\n')\n",
    "generated_passages = re.split(pattern, generated_text)[:-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_per = []\n",
    "count=0\n",
    "for l,p in zip(labels,generated_passages):\n",
    "  target_num = p.count(l)\n",
    "  all_num = sum([p.count(x) for x in all_label])\n",
    "  if all_num==0:\n",
    "    count+=1\n",
    "    continue\n",
    "  target_per.append(target_num/all_num)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.45526175371479094"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(target_per)/len(target_per)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing Prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/baihe/.conda/envs/tf36/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "from transformers import GPT2Tokenizer\n",
    "import spacy \n",
    "import re\n",
    "from tqdm import tnrange, tqdm_notebook\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model_name_or_path='gpt2',cache_dir=\"/data/baihe/dataset/gpt2\")\n",
    "nlp = spacy.blank(\"en\")\n",
    "nlp.add_pipe(nlp.create_pipe(\"sentencizer\"))\n",
    "def split_para(text):\n",
    "  text = text.replace(\"<newline> <newline>\",\"<newline>\")\n",
    "  paragraphs = text.split(\"<newline>\")\n",
    "  while \"\" in paragraphs:\n",
    "    paragraphs.remove(\"\")\n",
    "  return paragraphs\n",
    "def split_sent(text):\n",
    "  sentences = [x.text.strip() for x in nlp(text).sents]\n",
    "  return sentences\n",
    "def wp_preprocess(text):\n",
    "    # Standardize some symbols\n",
    "    text = text.replace('``', '\"')\n",
    "    text = text.replace(\"''\", '\"')\n",
    "    # Detokenize\n",
    "    text = re.sub(' +', ' ', text)\n",
    "    text = re.sub(' (\\'|\\.|\\,|\\:|\\?|\\!|;)', '\\g<1>', text)\n",
    "    text = re.sub('\" (.*) \"', '\"\\g<1>\"', text)\n",
    "    text = text.replace(\" n't\", \"n't\")\n",
    "    return text\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=272600), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "272600\n"
     ]
    }
   ],
   "source": [
    "prefix = \"/data/baihe/dataset/writingPrompts/original_data/\"\n",
    "data = [\"train\"]\n",
    "for name in data:\n",
    "  f_tar_out = open(prefix+name+\".wp_target_unfiltered\",'w',encoding='utf-8')\n",
    "  f_src_out = open(prefix+name+\".wp_source_unfiltered\",'w',encoding='utf-8')\n",
    "  with open(prefix+ name + \".wp_target\",'r',encoding='utf-8') as f_1:\n",
    "    with open(prefix+ name + \".wp_source\",'r',encoding='utf-8') as f_2:\n",
    "      src_all_data,tgt_all_data = f_2.readlines(),f_1.readlines()\n",
    "      total_examples=0\n",
    "      for source,target in tqdm_notebook(zip(src_all_data,tgt_all_data),total=len(src_all_data)):\n",
    "\n",
    "        src_passage = wp_preprocess(re.sub('\\[ (.*) \\]', '', source))\n",
    "        if src_passage=='':\n",
    "          continue\n",
    "        tgt_passage = wp_preprocess(target)\n",
    "        source_paras = split_para(src_passage.strip())\n",
    "        target_paras = split_para(tgt_passage.strip())\n",
    "#         text = \" \".join(source_paras+target_paras)\n",
    "#         total_tokens= len(source_paras)+len(target_paras)+3+len(tokenizer.tokenize(text))+len(split_sent(text))\n",
    "#         if total_tokens>1024:\n",
    "#           continue\n",
    "        total_examples+=1\n",
    "        \n",
    "        for para in source_paras:\n",
    "          sents = split_sent(para)\n",
    "          f_src_out.write(' [SENTENCE_EOS] '.join(sents)+' [SENTENCE_EOS]')\n",
    "          f_src_out.write(' [PARAGRAPH_EOS] ')\n",
    "        for para in target_paras:\n",
    "          sents = split_sent(para)\n",
    "          f_tar_out.write(' [SENTENCE_EOS] '.join(sents)+' [SENTENCE_EOS]')\n",
    "          f_tar_out.write(' [PARAGRAPH_EOS] ')\n",
    "        f_src_out.write('\\n')\n",
    "        f_tar_out.write('\\n')\n",
    "  print(total_examples)\n",
    "  f_tar_out.close()\n",
    "  f_src_out.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[ cw ] write about the strangest/scariest/saddest dream you 've ever had in less than 200 words . The night was as thick and terrifying as any I had ever seen before . All I could hear was the scream of the wind past my ears , the pounding of hooves , huffed horse breaths , and the pounding of my own heart .   The woods were closeknit , and my path was barely visible , hidden under a thick layer of bracken .   `` Faster , '' I whispered as I dug my heels in . Safety was close and yet so far away , calling to me . He would save me ; I knew it with all my heart .   All I had to do was outrun the demons at my back first .\""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=15620), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=272600), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=15138), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "prefix = \"/data/baihe/dataset/writingPrompts/original_data/\"\n",
    "data = [\"valid\",\"train\", \"test\"]\n",
    "for name in data:\n",
    "  f_tar_out = open(prefix+name+\".wp_target_filtered_w_newline\",'w',encoding='utf-8')\n",
    "  f_src_out = open(prefix+name+\".wp_source_filtered_w_newline\",'w',encoding='utf-8')\n",
    "  with open(prefix+ name + \".wp_target\",'r',encoding='utf-8') as f_1:\n",
    "    with open(prefix+ name + \".wp_source\",'r',encoding='utf-8') as f_2:\n",
    "      src_all_data,tgt_all_data = f_2.readlines(),f_1.readlines()\n",
    "      for source,target in tqdm_notebook(zip(src_all_data,tgt_all_data),total=len(src_all_data)):\n",
    "        if len(tokenizer.tokenize(source+target))>1024:\n",
    "          continue\n",
    "        f_src_out.write(source)\n",
    "        f_tar_out.write(target)\n",
    "  f_tar_out.close()\n",
    "  f_src_out.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=272600), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "30.029831254585474\n"
     ]
    }
   ],
   "source": [
    "prefix = \"/data/baihe/dataset/writingPrompts/original_data/\"\n",
    "name='train'\n",
    "with open(prefix+ name + \".wp_source\",'r',encoding='utf-8') as f_2:\n",
    "  src_all_data = f_2.readlines()\n",
    "  length = []\n",
    "  for source in tqdm_notebook(src_all_data,total=len(src_all_data)):\n",
    "    length.append(len(tokenizer.tokenize(source)))\n",
    "  print(sum(length)/len(length))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame({\"length\":length})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "272600"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "a=torch.tensor([1,2.3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.6500)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.mean(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}