draw_picture.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 画柱状图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置图片大小和像素\n",
    "plt.figure(figsize=(8,6),dpi=80)\n",
    "\n",
    "# setting the datas\n",
    "y1 = [531,\n",
    "541,\n",
    "507,\n",
    "489,\n",
    "514,\n",
    "500]\n",
    "y2 = [469,\n",
    "459,\n",
    "493,\n",
    "511,\n",
    "486,\n",
    "500]\n",
    "\n",
    "x = np.arange(len(y1))\n",
    "\n",
    "# 设置柱体宽度\n",
    "total_width, n = 0.7, 2\n",
    "width = total_width / n\n",
    "xticks_labels = [\"birth date\", \"birth place\", \"company\", \"major\", \"university\", \"workplace\"]\n",
    "yticks_labels = [0, 100, 200, 300, 400, 500, 600, 700]\n",
    "\n",
    "plt.barh(x - width/2, y1, height=width, fc=\"limegreen\", label=\"general\")\n",
    "plt.barh(x + width/2, y2, height=width, fc=\"aqua\", label=\"with spell/grammar error\")\n",
    "plt.yticks(x, xticks_labels)\n",
    "plt.xticks(yticks_labels, yticks_labels)\n",
    "# 设置标题\n",
    "plt.title(\"General vs Spell/Grammar error\",fontsize=16)\n",
    "\n",
    "# 设置坐标轴名称\n",
    "# plt.xlabel(\"features\")\n",
    "plt.xlabel(\"Sample size statistics\")\n",
    "\n",
    "# 设置图注\n",
    "plt.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置图片大小和像素\n",
    "plt.figure(figsize=(20,6),dpi=80)\n",
    "\n",
    "x = np.arange(6)\n",
    "\n",
    "# 设置柱体宽度\n",
    "total_width, n = 0.7, 2\n",
    "width = total_width / n\n",
    "style_labels = [\"Scientific_reports\", \"Novels\", \"Social_media\", \"Newspapers\"]\n",
    "xticks_labels = [\"birth date\", \"birth place\", \"company\", \"major\", \"university\", \"workplace\"]\n",
    "yticks_labels = [0, 100, 200, 300, 400, 500, 600, 700, 800]\n",
    "\n",
    "cur_sub_pic = 1\n",
    "prompt_type = \"statement\"\n",
    "for i in range(len(style_labels)):\n",
    "    for j in range(i + 1, len(style_labels)):\n",
    "        plt.subplot(2, 3, cur_sub_pic)\n",
    "        cur_sub_pic += 1\n",
    "\n",
    "        y1 = []\n",
    "        y2 = []\n",
    "        for feature in xticks_labels:\n",
    "            with open(\"./{}_vs_{}_ppl_acc_{}_{}.json\".format(style_labels[i], style_labels[j], feature.replace(\" \", \"_\"), prompt_type), 'r', encoding='utf8') as f:\n",
    "                cur_data = json.load(f)\n",
    "\n",
    "            y1.append(cur_data[\"choice {}\".format(style_labels[i])])\n",
    "            y2.append(cur_data[\"choice {}\".format(style_labels[j])])\n",
    "        \n",
    "        if sum(y1) < sum(y2):\n",
    "            plt.barh(x - width/2, y1, height=width, fc=\"limegreen\", label=style_labels[i])\n",
    "            plt.barh(x + width/2, y2, height=width, fc=\"aqua\", label=style_labels[j])\n",
    "            plt.yticks(x, xticks_labels)\n",
    "            plt.xticks(yticks_labels, yticks_labels)\n",
    "            # 设置图注\n",
    "            plt.legend()\n",
    "        else:\n",
    "            temp = y2\n",
    "            y2 = y1\n",
    "            y1 = temp\n",
    "\n",
    "            plt.barh(x - width/2, y1, height=width, fc=\"limegreen\", label=style_labels[j])\n",
    "            plt.barh(x + width/2, y2, height=width, fc=\"aqua\", label=style_labels[i])\n",
    "            plt.yticks(x, xticks_labels)\n",
    "            plt.xticks(yticks_labels, yticks_labels)\n",
    "            # 设置图注\n",
    "            plt.legend()\n",
    "            \n",
    "        # 设置坐标轴名称\n",
    "        plt.xlabel(\"Sample size statistics\")\n",
    "            \n",
    "# 设置标题\n",
    "# plt.title(\"General vs Spell/Grammar error\",fontsize=16)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 聚类"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TSNE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-01-30 23:19:49,328] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tqdm import tqdm\n",
    "from sklearn.manifold import TSNE\n",
    "from transformers import LlamaModel, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HIGH_DIM_VECS = {\n",
    "#     \"Blogs\": [],\n",
    "#     \"Personal Interviews\": [],\n",
    "#     \"Tabloids\": [],\n",
    "#     \"Textbooks\": [],\n",
    "#     \"Forum discussions\": [],\n",
    "#     \"Newspapers\": [],\n",
    "#     \"Social media\": [],\n",
    "#     \"Wikipedia\": [],\n",
    "#     \"Scientific reports\": [],\n",
    "#     \"Novels\": []\n",
    "# }\n",
    "# TYPE_LIST = [\"Scientific_reports\", \"Novels\", \"Forum_discussions\", \"Social_media\", \"Newspapers\", \"Wikipedia\", \"Blogs\", \"Personal_Interviews\", \"Textbooks\", \"Tabloids\"]\n",
    "HIGH_DIM_VECS = {\n",
    "    \"Newspapers\": [],\n",
    "    \"Social media\": [],\n",
    "    \"Scientific reports\": [],\n",
    "    \"Novels\": []\n",
    "}\n",
    "TYPE_LIST = [\"Scientific_reports\", \"Novels\", \"Social_media\", \"Newspapers\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first, let's see the birth_date information~\n",
    "def get_high_dim_vectors(model_name_or_path, test_file_path):\n",
    "    model = LlamaModel.from_pretrained(model_name_or_path, low_cpu_mem_usage=True, torch_dtype=torch.float16, device_map=\"auto\")\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
    "    model.eval()\n",
    "\n",
    "    with open(test_file_path, 'r', encoding='utf8') as f:\n",
    "        total_data = json.load(f)\n",
    "    for data_item in tqdm(total_data):\n",
    "        text_list = data_preprocess(data_item, tokenizer)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for text_item in text_list:\n",
    "                output = model(text_item[\"input_ids\"])\n",
    "                HIGH_DIM_VECS[text_item[\"style\"]].append(output.last_hidden_state[0][text_item[\"birthday_start_pos\"]:])\n",
    "    \n",
    "def data_preprocess(data_item, tokenizer):\n",
    "    first_birth_date = data_item[\"first_type_info\"][\"birth_date\"]\n",
    "    second_birth_date = data_item[\"second_type_info\"][\"birth_date\"]\n",
    "\n",
    "    text_list = []\n",
    "\n",
    "    for text in data_item[\"text_result\"]:\n",
    "        if first_birth_date in text:\n",
    "            key_word = first_birth_date\n",
    "            style = data_item[\"first_type_info\"][\"type_name\"]\n",
    "        elif second_birth_date in text:\n",
    "            key_word = second_birth_date\n",
    "            style = data_item[\"second_type_info\"][\"type_name\"]\n",
    "        else:\n",
    "            raise ValueError\n",
    "\n",
    "        pre_text = text.split(key_word)[0].strip()\n",
    "        post_text = pre_text + ' ' + key_word\n",
    "\n",
    "        pre_input_ids = tokenizer(pre_text, return_tensors='pt').input_ids\n",
    "        post_input_ids = tokenizer(post_text, return_tensors='pt').input_ids\n",
    "\n",
    "        input_ids = post_input_ids\n",
    "        start_pos = len(pre_input_ids[0])\n",
    "\n",
    "        text_list.append({\n",
    "            \"style\": style,\n",
    "            \"input_ids\": input_ids,\n",
    "            \"birthday_start_pos\": start_pos\n",
    "        })\n",
    "\n",
    "    return text_list\n",
    "\n",
    "def High_Dim_Vec_preprocess():\n",
    "    for key in list(HIGH_DIM_VECS.keys()):\n",
    "        HIGH_DIM_VECS[key] = torch.cat(HIGH_DIM_VECS[key], dim=0)\n",
    "\n",
    "def use_sne():\n",
    "    High_Dim_Vec_preprocess()\n",
    "    \n",
    "    input_vec = torch.cat([\n",
    "        # HIGH_DIM_VECS[\"Blogs\"], \n",
    "        # HIGH_DIM_VECS[\"Personal Interviews\"],\n",
    "        # HIGH_DIM_VECS[\"Tabloids\"],\n",
    "        # HIGH_DIM_VECS[\"Textbooks\"],\n",
    "        # HIGH_DIM_VECS[\"Forum discussions\"],\n",
    "        HIGH_DIM_VECS[\"Newspapers\"],\n",
    "        HIGH_DIM_VECS[\"Social media\"],\n",
    "        # HIGH_DIM_VECS[\"Wikipedia\"],\n",
    "        HIGH_DIM_VECS[\"Scientific reports\"],\n",
    "        HIGH_DIM_VECS[\"Novels\"],\n",
    "    ], dim=0)\n",
    "\n",
    "    labels = torch.cat([\n",
    "        # torch.ones(len(HIGH_DIM_VECS[\"Blogs\"])) * 1,\n",
    "        # torch.ones(len(HIGH_DIM_VECS[\"Personal Interviews\"])) * 2,\n",
    "        # torch.ones(len(HIGH_DIM_VECS[\"Tabloids\"])) * 3,\n",
    "        # torch.ones(len(HIGH_DIM_VECS[\"Textbooks\"])) * 4,\n",
    "        # torch.ones(len(HIGH_DIM_VECS[\"Forum discussions\"])) * 5,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Newspapers\"])) * 1,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Social media\"])) * 2,\n",
    "        # torch.ones(len(HIGH_DIM_VECS[\"Wikipedia\"])) * 8,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Scientific reports\"])) * 3,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Novels\"])) * 4,\n",
    "    ], dim=0)\n",
    "\n",
    "    tsne = TSNE(n_components=2, learning_rate=100, perplexity=10).fit_transform(input_vec)\n",
    "\n",
    "    plt.scatter(tsne[:, 0], tsne[:, 1], c=labels)\n",
    "\n",
    "    plt.colorbar()\n",
    "    plt.show()\n",
    "\n",
    "def main():\n",
    "    for i in range(len(TYPE_LIST)):\n",
    "        for j in range(i + 1, len(TYPE_LIST)):\n",
    "            type_i = TYPE_LIST[i]\n",
    "            type_j = TYPE_LIST[j]\n",
    "\n",
    "            model_path = \"/opt/tiger/fake_arnold/{}_vs_{}/checkpoint-780\".format(type_i, type_j)\n",
    "            data_path = \"./data_scripts/type_fights/bio_data_train_{}_vs_{}.json\".format(type_i, type_j)\n",
    "\n",
    "            get_high_dim_vectors(model_path, data_path)\n",
    "\n",
    "    # use_sne()\n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "Random = random.Random(666)\n",
    "\n",
    "SAMPLE_HIGH_DIM_VECS = {}\n",
    "\n",
    "for key in list(HIGH_DIM_VECS.keys()):\n",
    "    origin_list = HIGH_DIM_VECS[key]\n",
    "    Random.shuffle(origin_list)\n",
    "    SAMPLE_HIGH_DIM_VECS[key] = origin_list[:500]\n",
    "\n",
    "torch.save(SAMPLE_HIGH_DIM_VECS, \"temp_result.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['Newspapers', 'Social media', 'Scientific reports', 'Novels'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SAMPLE_HIGH_DIM_VECS.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in list(SAMPLE_HIGH_DIM_VECS.keys()):\n",
    "        SAMPLE_HIGH_DIM_VECS[key] = torch.cat(SAMPLE_HIGH_DIM_VECS[key], dim=0)\n",
    "\n",
    "\n",
    "input_vec = torch.cat([\n",
    "    SAMPLE_HIGH_DIM_VECS[\"Newspapers\"],\n",
    "    SAMPLE_HIGH_DIM_VECS[\"Social media\"],\n",
    "    SAMPLE_HIGH_DIM_VECS[\"Scientific reports\"],\n",
    "    SAMPLE_HIGH_DIM_VECS[\"Novels\"],\n",
    "], dim=0)\n",
    "\n",
    "labels = torch.cat([\n",
    "    torch.ones(len(SAMPLE_HIGH_DIM_VECS[\"Newspapers\"])) * 1,\n",
    "    torch.ones(len(SAMPLE_HIGH_DIM_VECS[\"Social media\"])) * 2,\n",
    "    torch.ones(len(SAMPLE_HIGH_DIM_VECS[\"Scientific reports\"])) * 3,\n",
    "    torch.ones(len(SAMPLE_HIGH_DIM_VECS[\"Novels\"])) * 4,\n",
    "], dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "tsne = TSNE(n_components=2, learning_rate=100, perplexity=30).fit_transform(input_vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(tsne[:, 0], tsne[:, 1], c=labels, s=2)\n",
    "\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first, let's see the birth_date information~\n",
    "def get_high_dim_vectors(model_name_or_path, test_file_path):\n",
    "    model = LlamaModel.from_pretrained(model_name_or_path, low_cpu_mem_usage=True, torch_dtype=torch.float16, device_map=\"auto\")\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
    "    model.eval()\n",
    "\n",
    "    with open(test_file_path, 'r', encoding='utf8') as f:\n",
    "        total_data = json.load(f)\n",
    "    for data_item in tqdm(total_data):\n",
    "        text_list = data_preprocess(data_item, tokenizer)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for text_item in text_list:\n",
    "                output = model(text_item[\"input_ids\"])\n",
    "                HIGH_DIM_VECS[text_item[\"style\"]].append(torch.mean(output.last_hidden_state[0], dim=0))\n",
    "    \n",
    "def data_preprocess(data_item, tokenizer):\n",
    "    text_list = []\n",
    "\n",
    "    for text in data_item[\"text_result\"]:\n",
    "        input_ids = tokenizer(text, return_tensors='pt').input_ids\n",
    "        style = \"None\"\n",
    "        for source in [\"Scientific_reports_info\", \"Novels_info\", \"Forum_discussions_info\", \"Social_media_info\", \"Newspapers_info\", \"Wikipedia_info\", \"Blogs_info\", \"Personal_Interviews_info\", \"Textbooks_info\", \"Tabloids_info\"]:\n",
    "            if data_item[source][\"birth_date\"] in text:\n",
    "                style = data_item[source][\"type_name\"]\n",
    "                break\n",
    "        \n",
    "        assert style != \"None\"\n",
    "\n",
    "        text_list.append({\n",
    "            \"style\": style,\n",
    "            \"input_ids\": input_ids,\n",
    "        })\n",
    "\n",
    "    return text_list\n",
    "\n",
    "def High_Dim_Vec_preprocess():\n",
    "    for key in list(HIGH_DIM_VECS.keys()):\n",
    "        HIGH_DIM_VECS[key] = torch.stack(HIGH_DIM_VECS[key], dim=0)\n",
    "        print(HIGH_DIM_VECS[key].shape)\n",
    "\n",
    "def use_sne():\n",
    "    High_Dim_Vec_preprocess()\n",
    "    \n",
    "    input_vec = torch.cat([\n",
    "        HIGH_DIM_VECS[\"Blogs\"], \n",
    "        HIGH_DIM_VECS[\"Personal Interviews\"],\n",
    "        HIGH_DIM_VECS[\"Tabloids\"],\n",
    "        HIGH_DIM_VECS[\"Textbooks\"],\n",
    "        HIGH_DIM_VECS[\"Forum discussions\"],\n",
    "        HIGH_DIM_VECS[\"Newspapers\"],\n",
    "        HIGH_DIM_VECS[\"Social media\"],\n",
    "        HIGH_DIM_VECS[\"Wikipedia\"],\n",
    "        HIGH_DIM_VECS[\"Scientific reports\"],\n",
    "        HIGH_DIM_VECS[\"Novels\"],\n",
    "    ], dim=0)\n",
    "\n",
    "    print(input_vec.shape)\n",
    "\n",
    "    labels = torch.cat([\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Blogs\"])) * 1,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Personal Interviews\"])) * 2,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Tabloids\"])) * 3,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Textbooks\"])) * 4,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Forum discussions\"])) * 5,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Newspapers\"])) * 6,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Social media\"])) * 7,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Wikipedia\"])) * 8,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Scientific reports\"])) * 9,\n",
    "        torch.ones(len(HIGH_DIM_VECS[\"Novels\"])) * 10,\n",
    "    ], dim=0)\n",
    "\n",
    "    pca = PCA(n_components=2)\n",
    "    x_pca = pca.fit_transform(input_vec)\n",
    "\n",
    "    plt.scatter(x_pca[:, 0], x_pca[:, 1], c=labels)\n",
    "\n",
    "    plt.colorbar()\n",
    "    plt.show()\n",
    "\n",
    "def main():\n",
    "    model_path = \"./all_styles_together/checkpoint-3905\"\n",
    "    data_path = \"./data_scripts/type_fights/bio_data_train_all_styles.json\"\n",
    "\n",
    "    get_high_dim_vectors(model_path, data_path)\n",
    "\n",
    "    use_sne()\n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def use_pca():\n",
    "    plt.figure(figsize=(15,10),dpi=80)\n",
    "    range = 1000\n",
    "    input_vec = torch.cat([\n",
    "        HIGH_DIM_VECS[\"Blogs\"][:range], \n",
    "        HIGH_DIM_VECS[\"Personal Interviews\"][:range],\n",
    "        HIGH_DIM_VECS[\"Tabloids\"][:range],\n",
    "        HIGH_DIM_VECS[\"Textbooks\"][:range],\n",
    "        HIGH_DIM_VECS[\"Forum discussions\"][:range],\n",
    "        HIGH_DIM_VECS[\"Newspapers\"][:range],\n",
    "        HIGH_DIM_VECS[\"Social media\"][:range],\n",
    "        HIGH_DIM_VECS[\"Wikipedia\"][:range],\n",
    "        HIGH_DIM_VECS[\"Scientific reports\"][:range],\n",
    "        HIGH_DIM_VECS[\"Novels\"][:range],\n",
    "    ], dim=0)\n",
    "\n",
    "    print(input_vec.shape)\n",
    "\n",
    "    pca = PCA(n_components=2)\n",
    "    x_pca = pca.fit_transform(input_vec)\n",
    "\n",
    "    plt.scatter(x_pca[:, 0][:range], x_pca[:, 1][:range], label=\"Blogs\", c=\"lightseagreen\", s=1)\n",
    "    plt.scatter(x_pca[:, 0][6 * range:7 * range], x_pca[:, 1][6 * range:7 * range], label=\"Social media\", c='b', s=1)\n",
    "    plt.scatter(x_pca[:, 0][4 * range:5 * range], x_pca[:, 1][4 * range:5 * range], label=\"Forum discussions\", c='deepskyblue', s=1)\n",
    "    plt.scatter(x_pca[:, 0][2 * range:3 * range], x_pca[:, 1][2 * range:3 * range], label=\"Tabloids\",c='slateblue', s=1)\n",
    "\n",
    "    plt.scatter(x_pca[:, 0][3 * range:4 * range], x_pca[:, 1][3 * range:4 * range], label=\"Textbooks\", c='r', s=1)\n",
    "    plt.scatter(x_pca[:, 0][7 * range:8 * range], x_pca[:, 1][7 * range:8 * range], label=\"Wikipedia\", c=\"darkred\", s=1)\n",
    "    plt.scatter(x_pca[:, 0][8 * range:9 * range], x_pca[:, 1][8 * range:9 * range], label=\"Scientific reports\",c=\"gold\", s=1)    \n",
    "    plt.scatter(x_pca[:, 0][5 * range:6 * range], x_pca[:, 1][5 * range:6 * range], label=\"Newspapers\", c='orange', s=1)\n",
    "    \n",
    "    plt.scatter(x_pca[:, 0][9 * range:], x_pca[:, 1][9 * range:], label=\"Novels\", c='green', s=1)\n",
    "    plt.scatter(x_pca[:, 0][range:2 * range], x_pca[:, 1][range: 2 * range], label=\"Personal Interviews\", c='gray', s=1)\n",
    "\n",
    "    plt.legend(fontsize=8)\n",
    "    plt.show()\n",
    "\n",
    "use_pca()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 折线图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_size = [\"14m\", \"70m\", \"160m\", \"410m\", \"1b\", \"1.4b\", \"2.8b\", \"6.9b\"]\n",
    "x = range(0, 8)\n",
    "y = [450, 500, 550, 600, 650, 700, 750, 800, 900, 1000]\n",
    "y_labels = [\"45%\", \"50%\", \"55%\", \"60%\", \"65%\", \"70%\", \"75%\", \"80%\", \"90%\", \"100%\"]\n",
    "# 先只写newspaper的比例\n",
    "prompt_type=\"statement\"\n",
    "\n",
    "data_collect = {\n",
    "\"birth_date\" : [],\n",
    "\"birth_place\" : [],\n",
    "\"company\" : [],\n",
    "\"major\" : [],\n",
    "\"university\" : [],\n",
    "\"workplace\" : [],\n",
    "\"avg\" : [],\n",
    "}\n",
    "\n",
    "for feature_type in [\"birth_date\", \"birth_place\", \"company\", \"major\", \"university\", \"workplace\"]:\n",
    "    for size in model_size:\n",
    "        with open(\"./pythia_{}_Social_media_vs_Newspapers_ppl_acc_{}_{}.json\".format(size, feature_type, prompt_type), 'r', encoding='utf8') as f:\n",
    "            result = json.load(f)\n",
    "\n",
    "        data_collect[feature_type].append(result[\"choice Newspapers\"])\n",
    "\n",
    "plt.figure(figsize=(9,5.562),dpi=80)\n",
    "plt.title(\"Variation of model preferences with model scale\")\n",
    "plt.xticks(x, model_size)\n",
    "plt.yticks(y, y_labels)\n",
    "\n",
    "plt.plot(x, data_collect[\"birth_date\"], label='birth date')\n",
    "plt.plot(x, data_collect[\"birth_place\"], label='birth place')\n",
    "plt.plot(x, data_collect[\"company\"], label='company')\n",
    "plt.plot(x, data_collect[\"major\"], label='major')\n",
    "plt.plot(x, data_collect[\"university\"], label='university')\n",
    "plt.plot(x, data_collect[\"workplace\"], label='work place')\n",
    "# plt.plot(x, avg, label=\"avg\")\n",
    "\n",
    "plt.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_list = [\"paper_num\", \"sat_score\", \"birth_date\"]\n",
    "prompt_list = [\"question\", \"statement\"]\n",
    "precentage_list = [\"68%\", \"70%\", \"72%\", \"74%\", \"76%\", \"78%\", \"80%\", \"82%\", \"84%\"]\n",
    "x = range(3)\n",
    "y = [680, 700, 720, 740, 760, 780, 800, 820, 840]\n",
    "\n",
    "for prompt_type in prompt_list:\n",
    "    results = []\n",
    "    for feature in feature_list:\n",
    "        with open(\"checkpoint-785_ppl_acc_{}_{}.json\".format(feature, prompt_type), 'r', encoding='utf8') as f:\n",
    "            cur_data = json.load(f)\n",
    "        results.append(cur_data[\"choice Newspapers\"])\n",
    "    plt.plot(x, results, label=\"{} result\".format(prompt_type))\n",
    "    # plt.bar(x, results, label=\"{} result\".format(prompt_type), bottom=65)\n",
    "\n",
    "plt.xticks(x, feature_list)\n",
    "plt.yticks(y, precentage_list)\n",
    "plt.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 饼状图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置图片大小和像素\n",
    "plt.figure(figsize=(8,6),dpi=80)\n",
    "\n",
    "x_labels = [\"birth date\", \"birth place\", \"company\", \"major\", \"university\"]\n",
    "styles = [\"Scientific_reports\", \"Wikipedia\", \"Newspapers\", \"Textbooks\", \"Novels\", \"Personal_Interviews\", \"Forum_discussions\", \"Social_media\",  \"Blogs\", \"Tabloids\"]\n",
    "colors = [\"gold\", \"darkred\", \"orange\", \"r\", \"green\", \"gray\", \"deepskyblue\", \"b\", \"lightseagreen\", \"slateblue\"]\n",
    "x = np.arange(len(x_labels))\n",
    "y = [0, 200, 400, 600, 800, 1000]\n",
    "y_labels = [\"0%\", \"20%\", \"40%\", \"60%\", \"80%\", \"100%\"]\n",
    "# bottom = [0] * len(x_labels)\n",
    "nums = []\n",
    "\n",
    "for sty_idx, style in enumerate(styles):\n",
    "    style_result = []\n",
    "    for feature in x_labels:\n",
    "        with open(\"all_styles_together_ppl_acc_{}_statement.json\".format(feature.replace(\" \", \"_\")), 'r', encoding='utf8') as f:\n",
    "            total_data = json.load(f)\n",
    "        style_result.append(total_data[\"result\"][style][0])\n",
    "\n",
    "    nums.append(sum(style_result))\n",
    "\n",
    "    # plt.bar(x, style_result, width=0.5, bottom=bottom, fc=colors[sty_idx], label=style.replace(\"_\", \" \"))\n",
    "    # bottom = list(map(lambda x, y: x + y, style_result, bottom))\n",
    "plt.pie(nums, labels=styles, explode=[0.02] * 10)\n",
    "plt.legend(bbox_to_anchor=(1.05, 0), loc=3, borderaxespad=0)\n",
    "\n",
    "# plt.xticks(x, x_labels)\n",
    "# plt.yticks(y, y_labels)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.2 64-bit",
   "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.9.2"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}