knowledge_finding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.decomposition import PCA\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "from transformers import BitsAndBytesConfig\n",
    "from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hidden_states(model, input, tokenizer, device):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        input = tokenizer(input, return_tensors='pt')\n",
    "        output = model(input.input_ids.to(device), attention_mask=input.attention_mask.to(device), output_hidden_states=True)\n",
    "\n",
    "    # Tuple of  (batch_size, sequence_length, hidden_size)\n",
    "    return output.hidden_states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_all_hidden_states(model, dataset:list, tokenizer, device):\n",
    "    result = {}\n",
    "\n",
    "    for data in tqdm(dataset):\n",
    "        hidden_states = get_hidden_states(model, data, tokenizer, device)\n",
    "\n",
    "        layer_num = model.config.num_hidden_layers\n",
    "        for layer_id in range(layer_num + 1):\n",
    "            vector_in_tensor = hidden_states[layer_id][0].to(\"cpu\")\n",
    "            vector_in_tensor = torch.mean(vector_in_tensor, dim=0)\n",
    "            try:\n",
    "                result[layer_id].append(vector_in_tensor)\n",
    "            except:\n",
    "                result[layer_id] = [vector_in_tensor]\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, json_path):\n",
    "        with open(json_path, 'r') as f:\n",
    "            self.data = json.load(f)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        # 根据实际情况处理数据\n",
    "        sample = self.data[idx]\n",
    "        return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_prefix = {'en': '''It starts off a bit slow, but once the product placement jokes start it takes off.\n",
    "Emotion: \"positive\"             \n",
    "I've read this book with much expectation, it was very boring all through out the book\n",
    "Emotion: \"negative\"''',\n",
    "                 'zh': '''一开始有点慢，但一旦植入式笑话开始，它就会起飞。\n",
    "情感: \"积极\"\n",
    "我带着很大的期待读了这本书，整本书都很无聊\n",
    "情感: \"消极\"''',\n",
    "                  'hi': '''इसकी शुरुआत थोड़ी धीमी होती है, लेकिन एक बार जब उत्पाद प्लेसमेंट के बारे में मजाक शुरू हो जाता है तो इसमें तेजी आ जाती है।\n",
    "भावना: \"सकारात्मक\"\n",
    "मैंने यह किताब बहुत उम्मीद के साथ पढ़ी है, पूरी किताब में यह बहुत उबाऊ थी\n",
    "भावना: \"नकारात्मक\"''',\n",
    "                  'de':'''Es fängt etwas langsam an, aber sobald die Witze über die Produktplatzierung beginnen, geht es los.\n",
    "Emotion: \"positiv\"\n",
    "Ich habe dieses Buch mit großer Erwartung gelesen, es war durchweg sehr langweilig\n",
    "Emotion: \"negativ\"''',\n",
    "                  'fr':'''Cela commence un peu lentement, mais une fois que les blagues sur le placement de produit commencent, cela décolle.\n",
    "Émotion: \"positif\"\n",
    "J'ai lu ce livre avec beaucoup d'attente, c'était très ennuyeux tout au long du livre\n",
    "Émotion: \"négatif\"''',\n",
    "                  'th':'''มันเริ่มต้นช้านิดหน่อย แต่เมื่อเรื่องตลกเกี่ยวกับการจัดวางผลิตภัณฑ์เริ่มต้นขึ้น\n",
    "อารมณ์: \"เชิงบวก\"\n",
    "ฉันอ่านหนังสือเล่มนี้ด้วยความคาดหวังมาก มันน่าเบื่อมากตลอดทั้งเล่ม\n",
    "อารมณ์: \"เชิงลบ\"''',\n",
    "                  'sw':'''Huanza polepole, lakini utani wa uwekaji bidhaa unapoanza huanza.\n",
    "Hisia: \"chanya\"\n",
    "Nimekisoma kitabu hiki kwa matarajio mengi, kilikuwa cha kuchosha katika kitabu chote\n",
    "Hisia: \"hasi\"''',\n",
    "                  'ms':'''Ia bermula agak perlahan, tetapi apabila jenaka peletakan produk bermula, ia bermula.\n",
    "Emosi: \"positif\"\n",
    "Saya telah membaca buku ini dengan penuh harapan, ia sangat membosankan sepanjang buku itu\n",
    "Emosi: \"negatif\"''',}\n",
    "\n",
    "emotion_map = {'en':'Emotion: ', 'zh':'情感: ', 'hi':'भावना: ', 'de':'Emotion: ', 'fr':'Émotion: ',\n",
    "               'th':'อารมณ์: ', 'sw':'Hisia: ', 'ms':'Emosi: '}\n",
    "\n",
    "\n",
    "def preprocess_data(data_path, language):\n",
    "    data_list = []\n",
    "    dataset = CustomDataset(data_path)\n",
    "    assert(dataset is not None)\n",
    "    dataloader = DataLoader(dataset, batch_size=1, shuffle=False)\n",
    "    sample_iterator = tqdm(dataloader, desc=\"Inferenced batchs num\")\n",
    "    for step, inputs in enumerate(sample_iterator):\n",
    "        text_item = prompt_prefix[language] + '\\n' + inputs['input'][0] + '\\n' + emotion_map[language] +  '\\\"'\n",
    "        data_list.append(text_item)\n",
    "\n",
    "    return data_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pca_decomp(data:list, n_components=2, if_save=False, save_dir=None, model_size=None, layer=0):\n",
    "    \n",
    "    # labels = torch.cat([\n",
    "    #     torch.ones(100) * 1,\n",
    "    #     torch.ones(100) * 2,\n",
    "    #     torch.ones(100) * 3,\n",
    "    #     torch.ones(100) * 4,\n",
    "    #     # torch.ones(100) * 5,\n",
    "    #     # torch.ones(100) * 6,\n",
    "    #     # torch.ones(100) * 7,\n",
    "    #     # torch.ones(100) * 8,\n",
    "    #     # torch.ones(100) * 9,\n",
    "    #     # torch.ones(100) * 10,\n",
    "    # ], dim=0)\n",
    "\n",
    "    input_vec = torch.stack(tensors=data, dim=0)\n",
    "    pca = PCA(n_components=n_components)\n",
    "\n",
    "    pca_result = pca.fit_transform(input_vec.numpy())\n",
    "\n",
    "    plt.cla()\n",
    "\n",
    "\n",
    "    # plt.colorbar()\n",
    "\n",
    "    if if_save:\n",
    "        plt.scatter(pca_result[:500, 0], pca_result[:500, 1], label=\"en\")\n",
    "        plt.scatter(pca_result[500:1000, 0], pca_result[500:1000, 1], label=\"de\")\n",
    "        plt.scatter(pca_result[1000:1500, 0], pca_result[1000:1500, 1], label=\"fr\")\n",
    "        plt.scatter(pca_result[1500:, 0], pca_result[1500:, 1], label=\"hi\")\n",
    "\n",
    "        plt.legend(fontsize=16)\n",
    "        os.makedirs(f'{os.path.join(save_dir, model_size)}/PCA/emotion/2-shot/figures', exist_ok=True)\n",
    "        plt.savefig(f'{os.path.join(save_dir, model_size)}/PCA/emotion/2-shot/figures/{model_size}_pca_{layer}.pdf', dpi=300, bbox_inches='tight')\n",
    "    else:\n",
    "        plt.scatter(pca_result[:500, 0], pca_result[:500, 1], label=\"en\")\n",
    "        plt.scatter(pca_result[500:1000, 0], pca_result[500:1000, 1], label=\"de\")\n",
    "        plt.scatter(pca_result[1000:1500, 0], pca_result[1000:1500, 1], label=\"fr\")\n",
    "        plt.scatter(pca_result[1500:, 0], pca_result[1500:, 1], label=\"hi\")\n",
    "\n",
    "        plt.legend(fontsize=16)\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pca_decomp_seperate(data:list, n_components=2, if_save=False, save_dir=None, model_size=None, layer=0):\n",
    "    \n",
    "    # labels = torch.cat([\n",
    "    #     torch.ones(100) * 1,\n",
    "    #     torch.ones(100) * 2,\n",
    "    #     torch.ones(100) * 3,\n",
    "    #     torch.ones(100) * 4,\n",
    "    #     # torch.ones(100) * 5,\n",
    "    #     # torch.ones(100) * 6,\n",
    "    #     # torch.ones(100) * 7,\n",
    "    #     # torch.ones(100) * 8,\n",
    "    #     # torch.ones(100) * 9,\n",
    "    #     # torch.ones(100) * 10,\n",
    "    # ], dim=0)\n",
    "\n",
    "    input_vec = torch.stack(tensors=data, dim=0)\n",
    "    pca = PCA(n_components=n_components)\n",
    "\n",
    "    pca_result = pca.fit_transform(input_vec.numpy())\n",
    "\n",
    "    plt.cla()\n",
    "\n",
    "    # print((pca_result[:10, 0]).type)\n",
    "\n",
    "\n",
    "    # plt.colorbar()\n",
    "\n",
    "    if if_save:\n",
    "        plt.scatter(np.concatenate((pca_result[:100*1, 0], pca_result[500:500+100*1, 0], pca_result[1000:1000+100*1, 0], pca_result[1500:1500+100*1, 0])), np.concatenate((pca_result[:100*1, 1], pca_result[500:500+100*1, 1], pca_result[1000:1000+100*1, 1], pca_result[1500:1500+100*1, 1])), label=\"0-100\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*1:100*2, 0], pca_result[500+100*1:500+100*2, 0], pca_result[1000+100*1:1000+100*2, 0], pca_result[1500+100*1:1500+100*2, 0])), np.concatenate((pca_result[100*1:100*2, 1], pca_result[500+100*1:500+100*2, 1], pca_result[1000+100*1:1000+100*2, 1], pca_result[1500+100*1:1500+100*2, 1])), label=\"100-200\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*2:100*3, 0], pca_result[500+100*2:500+100*3, 0], pca_result[1000+100*2:1000+100*3, 0], pca_result[1500+100*2:1500+100*3, 0])), np.concatenate((pca_result[100*2:100*3, 1], pca_result[500+100*2:500+100*3, 1], pca_result[1000+100*2:1000+100*3, 1], pca_result[1500+100*2:1500+100*3, 1])), label=\"200-300\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*3:100*4, 0], pca_result[500+100*3:500+100*4, 0], pca_result[1000+100*3:1000+100*4, 0], pca_result[1500+100*3:1500+100*4, 0])), np.concatenate((pca_result[100*3:100*4, 1], pca_result[500+100*3:500+100*4, 1], pca_result[1000+100*3:1000+100*4, 1], pca_result[1500+100*3:1500+100*4, 1])), label=\"300-400\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*4:100*5, 0], pca_result[500+100*4:500+100*5, 0], pca_result[1000+100*4:1000+100*5, 0], pca_result[1500+100*4:1500+100*5, 0])), np.concatenate((pca_result[100*4:100*5, 1], pca_result[500+100*4:500+100*5, 1], pca_result[1000+100*4:1000+100*5, 1], pca_result[1500+100*4:1500+100*5, 1])), label=\"400-500\")\n",
    "\n",
    "        plt.legend(fontsize=16)\n",
    "        os.makedirs(f'{os.path.join(save_dir, model_size)}/PCA/emotion/2-shot/figures', exist_ok=True)\n",
    "        plt.savefig(f'{os.path.join(save_dir, model_size)}/PCA/emotion/2-shot/figures/{model_size}_pca_{layer}.pdf', dpi=300, bbox_inches='tight')\n",
    "    else:\n",
    "        plt.scatter(np.concatenate((pca_result[:100*1, 0], pca_result[500:500+100*1, 0], pca_result[1000:1000+100*1, 0], pca_result[1500:1500+100*1, 0])), np.concatenate((pca_result[:100*1, 1], pca_result[500:500+100*1, 1], pca_result[1000:1000+100*1, 1], pca_result[1500:1500+100*1, 1])), label=\"0-100\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*1:100*2, 0], pca_result[500+100*1:500+100*2, 0], pca_result[1000+100*1:1000+100*2, 0], pca_result[1500+100*1:1500+100*2, 0])), np.concatenate((pca_result[100*1:100*2, 1], pca_result[500+100*1:500+100*2, 1], pca_result[1000+100*1:1000+100*2, 1], pca_result[1500+100*1:1500+100*2, 1])), label=\"100-200\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*2:100*3, 0], pca_result[500+100*2:500+100*3, 0], pca_result[1000+100*2:1000+100*3, 0], pca_result[1500+100*2:1500+100*3, 0])), np.concatenate((pca_result[100*2:100*3, 1], pca_result[500+100*2:500+100*3, 1], pca_result[1000+100*2:1000+100*3, 1], pca_result[1500+100*2:1500+100*3, 1])), label=\"200-300\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*3:100*4, 0], pca_result[500+100*3:500+100*4, 0], pca_result[1000+100*3:1000+100*4, 0], pca_result[1500+100*3:1500+100*4, 0])), np.concatenate((pca_result[100*3:100*4, 1], pca_result[500+100*3:500+100*4, 1], pca_result[1000+100*3:1000+100*4, 1], pca_result[1500+100*3:1500+100*4, 1])), label=\"300-400\")\n",
    "        plt.scatter(np.concatenate((pca_result[100*4:100*5, 0], pca_result[500+100*4:500+100*5, 0], pca_result[1000+100*4:1000+100*5, 0], pca_result[1500+100*4:1500+100*5, 0])), np.concatenate((pca_result[100*4:100*5, 1], pca_result[500+100*4:500+100*5, 1], pca_result[1000+100*4:1000+100*5, 1], pca_result[1500+100*4:1500+100*5, 1])), label=\"400-500\")\n",
    "\n",
    "        plt.legend(fontsize=16)\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pca_decomp_pearson(data:list, n_components=1, model_size=None, layer=0):\n",
    "    \n",
    "    print(f\"Layer is {layer}\")\n",
    "\n",
    "    # labels = torch.cat([\n",
    "    #     torch.ones(100) * 1,\n",
    "    #     torch.ones(100) * 2,\n",
    "    #     torch.ones(100) * 3,\n",
    "    #     torch.ones(100) * 4,\n",
    "    #     # torch.ones(100) * 5,\n",
    "    #     # torch.ones(100) * 6,\n",
    "    #     # torch.ones(100) * 7,\n",
    "    #     # torch.ones(100) * 8,\n",
    "    #     # torch.ones(100) * 9,\n",
    "    #     # torch.ones(100) * 10,\n",
    "    # ], dim=0)\n",
    "\n",
    "    input_vec = torch.stack(tensors=data, dim=0)\n",
    "    pca = PCA(n_components=n_components)\n",
    "\n",
    "    pca_result = pca.fit_transform(input_vec.numpy())\n",
    "\n",
    "    plt.cla()\n",
    "\n",
    "    # print((pca_result[:10, 0]).type)\n",
    "\n",
    "\n",
    "    # plt.colorbar()\n",
    "\n",
    "    en = pca_result[:500, 0]\n",
    "    de = pca_result[500:1000, 0]\n",
    "    fr = pca_result[1000:1500, 0]\n",
    "    hi = pca_result[1500:2000, 0]\n",
    "    th = pca_result[2000:2500, 0]\n",
    "    sw = pca_result[2500:3000, 0]\n",
    "    ms = pca_result[3000:, 0]\n",
    "\n",
    "    pearson_ende = np.corrcoef(en, de)[0, 1]\n",
    "    pearson_enfr = np.corrcoef(en, fr)[0, 1]\n",
    "    pearson_enhi = np.corrcoef(en, hi)[0, 1]\n",
    "    pearson_defr = np.corrcoef(de, fr)[0, 1]\n",
    "    pearson_dehi = np.corrcoef(de, hi)[0, 1]\n",
    "    pearson_frhi = np.corrcoef(fr, hi)[0, 1]\n",
    "    pearson_enth = np.corrcoef(en, th)[0, 1]\n",
    "    pearson_ensw = np.corrcoef(en, sw)[0, 1]\n",
    "    pearson_enms = np.corrcoef(en, ms)[0, 1]\n",
    "    \n",
    "    print(f\"Model size: {model_size}\")\n",
    "    print(f\"pearson_ende: {pearson_ende}\")\n",
    "    print(f\"pearson_enfr: {pearson_enfr}\")\n",
    "    print(f\"pearson_enhi: {pearson_enhi}\")\n",
    "    print(f\"pearson_defr: {pearson_defr}\")\n",
    "    print(f\"pearson_dehi: {pearson_dehi}\")\n",
    "    print(f\"pearson_frhi: {pearson_frhi}\")\n",
    "    print(f\"pearson_enth: {pearson_enth}\")\n",
    "    print(f\"pearson_ensw: {pearson_ensw}\")\n",
    "    print(f\"pearson_enms: {pearson_enms}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_size = 'Qwen1.5-1.8b-aligned'\n",
    "\n",
    "all_avaliable_model = {'mistral-7b':'/home/nfs02/model/mistralai_Mistral-7B-v0.1',\n",
    "                       'mistral-7b-aligned':'/home/nfs03/zhangsm/multiL-transfer-interpretability/pretrained-models/mistral_zhit20k_round1_epoch3',\n",
    "                       'Qwen1.5-0.5b':'/home/nfs02/model/Qwen1.5-0.5B',\n",
    "                       'Qwen1.5-1.8b':'/home/nfs02/model/Qwen1.5-1.8B',\n",
    "                       'Qwen1.5-1.8b-aligned':'/home/nfs03/zhangsm/multiL-transfer-interpretability/pretrained-models/Qwen1.8b_emotion_swhi20k_round1_epoch3',\n",
    "                       'Qwen1.5-4b':'/home/nfs02/model/Qwen1.5-4B',\n",
    "                       'Qwen1.5-4b-aligned':'/home/nfs03/zhangsm/multiL-transfer-interpretability/pretrained-models/Qwen4b_emotion_swhi20k_round1_epoch3',\n",
    "                       'Qwen1.5-14b':'/home/nfs02/model/Qwen1.5-14B-Base',\n",
    "                       'Qwen1.5-14b-aligned':'/home/nfs03/zhangsm/multiL-transfer-interpretability/pretrained-models/Qwen14b_emotion_swhi20k_round1_epoch3'}\n",
    "model_name_or_path = all_avaliable_model[model_size]\n",
    "\n",
    "quantization_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.float16\n",
    ")\n",
    "\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_auth_token=None)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name_or_path, use_auth_token=None, device_map='auto', quantization_config=quantization_config)\n",
    "# model.to(device)\n",
    "hidden_states = {}\n",
    "\n",
    "\n",
    "for language in [\"en\", \"de\", \"fr\", \"hi\", \"th\", \"sw\", \"ms\"]:\n",
    "    data_path = \"/home/nfs03/zhangsm/multiL-transfer-interpretability/zhangsm-multiL/data/ap_emotion/emotion_{}500_test.json\".format(language)\n",
    "    data_list = preprocess_data(data_path, language)\n",
    "\n",
    "    hidden_states[language] = get_all_hidden_states(model, data_list, tokenizer, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get pca result\n",
    "# layer = 32\n",
    "# input_tensor = [\n",
    "#     item for lang in [\"en\", \"de\", \"fr\", \"hi\"] for item in hidden_states[lang][layer][:500]\n",
    "# ]\n",
    "\n",
    "# pca_decomp(input_tensor, if_save=False, save_dir='../zsm-results', model_size=model_size, layer=layer)\n",
    "\n",
    "\n",
    "# for layer in range(0, 25):\n",
    "#     input_tensor = [\n",
    "#         item for lang in [\"en\", \"de\", \"fr\", \"hi\"] for item in hidden_states[lang][layer][:500]\n",
    "#     ]\n",
    "\n",
    "#     pca_decomp(input_tensor, if_save=True, save_dir='./zsm-results', model_size=model_size, layer=layer)\n",
    "\n",
    "\n",
    "# layer = 20\n",
    "# input_tensor = [\n",
    "#     item for lang in [\"en\", \"de\", \"fr\", \"hi\"] for item in hidden_states[lang][layer][:500]\n",
    "# ]\n",
    "\n",
    "# pca_decomp_seperate(input_tensor, if_save=False, save_dir='../zsm-results', model_size=model_size, layer=layer)\n",
    "\n",
    "\n",
    "layer = 18\n",
    "input_tensor = [\n",
    "    item for lang in [\"en\", \"de\", \"fr\", \"hi\", \"th\", \"sw\", \"ms\"] for item in hidden_states[lang][layer][:500]\n",
    "]\n",
    "\n",
    "pca_decomp_pearson(input_tensor, n_components=1, model_size=model_size, layer=layer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "caoyq_alpaca_lora",
   "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}