sentence_transformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51fa7f9d-5dc1-45fa-a4c8-00fce7fcdb53",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd \"/data/KunZhao/on-isotropy-of-contrastive-SRL-main/dd++\"\n",
    "# magic code to automatically reload external modules if they change\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf381e50-bb2a-4ea9-a874-313882694a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import sys\n",
    "from torch.utils.data import DataLoader\n",
    "from sentence_transformers import SentenceTransformer, LoggingHandler, util, models, evaluation, losses, InputExample\n",
    "import logging\n",
    "from datetime import datetime\n",
    "import gzip\n",
    "import os\n",
    "import tarfile\n",
    "from collections import defaultdict\n",
    "from torch.utils.data import IterableDataset\n",
    "import tqdm\n",
    "from torch.utils.data import Dataset\n",
    "import random\n",
    "import pickle\n",
    "import argparse\n",
    "from sentence_transformers import evaluation\n",
    "from sentence_transformers import InputExample\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b490ac6c-4ef9-40fd-bfb5-86e325f50bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!export CUDA_VISIBLE_DEVICES=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56bd8949-e654-48e2-b464-104f3510a50c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers.models.bert.modeling_bert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee01eddb-a83a-4c47-8ced-66b4ba4078a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "229c3c0d-fdef-4c0e-b9df-60b74bd47980",
   "metadata": {},
   "outputs": [],
   "source": [
    "!conda list\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ae64bc8-41f7-44f5-93f1-d9c56c92ece2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"train.json\", \"r\") as f:\n",
    "    lines = f.readlines()\n",
    "print(type(lines))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "447c926f-e4f8-4c19-bacc-5e5903db66b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(lines))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbbde817-3bd8-42b7-926e-04618f3cfc3c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# lines_len = []\n",
    "train_examples = []\n",
    "for line in lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            train_examples.append([context, pos, neg])\n",
    "print(sample)\n",
    "print(train_examples[-1])\n",
    "                            \n",
    "    # lines_len.append(len(ss.strip()))\n",
    "# print(\"context: \")\n",
    "# print(sample[\"context\"])\n",
    "# print(\"\\n\")\n",
    "# print(\"positive_responses:\")\n",
    "# for pos in sample[\"positive_responses\"]:\n",
    "#     print(pos)\n",
    "\n",
    "# print(\"\\n\")\n",
    "# print(\"adversarial_negative_responses:\")\n",
    "# for adv in sample[\"adversarial_negative_responses\"]:\n",
    "#     print(adv)\n",
    "\n",
    "# print(\"\\n\")\n",
    "# print(\"random_negative_responses:\")\n",
    "# for ran in sample[\"random_negative_responses\"]:\n",
    "#     print(ran)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35a9158d-4053-4698-a93a-6f4d2cd230e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = 0\n",
    "\n",
    "for l in lines_len:\n",
    "    if l > 512:\n",
    "       c = c + 1 \n",
    "print(c, len(lines_len))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2551e785-b95c-4279-864d-d2fc3f973c08",
   "metadata": {},
   "outputs": [],
   "source": [
    "ss = \"[CLS] \"\n",
    "for c in sample[\"context\"]:\n",
    "    ss = ss + c + \" [SEP] \"\n",
    "print(ss.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbc44c34-3cfa-42e2-884a-ca77c8f901a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_context = "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a063f804-6ba1-4769-90a7-8b9654714721",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import InputExample\n",
    "\n",
    "train_examples = []\n",
    "for line in lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            train_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "# print(sample)\n",
    "# print(train_examples[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58a81912-b3bb-4d79-b3fc-3d7e991b90c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(train_examples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7caa9a3b-c90a-4917-8242-20c912964b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_examples[0])\n",
    "print(train_examples[1])\n",
    "print(train_examples[2])\n",
    "print(train_examples[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad2bcc5e-0b41-42b6-ad94-80ce9ecc5be6",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This examples show how to train a Bi-Encoder for the MS Marco dataset (https://github.com/microsoft/MSMARCO-Passage-Ranking).\n",
    "\n",
    "The queries and passages are passed independently to the transformer network to produce fixed sized embeddings.\n",
    "These embeddings can then be compared using cosine-similarity to find matching passages for a given query.\n",
    "\n",
    "For training, we use MultipleNegativesRankingLoss. There, we pass triplets in the format:\n",
    "(query, positive_passage, negative_passage)\n",
    "\n",
    "Negative passage are hard negative examples, that were mined using different dense embedding methods and lexical search methods.\n",
    "Each positive and negative passage comes with a score from a Cross-Encoder. This allows denoising, i.e. removing false negative\n",
    "passages that are actually relevant for the query.\n",
    "\n",
    "With a distilbert-base-uncased model, it should achieve a performance of about 33.79 MRR@10 on the MSMARCO Passages Dev-Corpus\n",
    "\n",
    "Running this script:\n",
    "python train_bi-encoder-v3.py\n",
    "\"\"\"\n",
    "import sys\n",
    "import json\n",
    "from torch.utils.data import DataLoader\n",
    "from sentence_transformers import SentenceTransformer, LoggingHandler, util, models, evaluation, losses, InputExample\n",
    "import logging\n",
    "from datetime import datetime\n",
    "import gzip\n",
    "import os\n",
    "import tarfile\n",
    "from collections import defaultdict\n",
    "from torch.utils.data import IterableDataset\n",
    "import tqdm\n",
    "from torch.utils.data import Dataset\n",
    "import random\n",
    "import pickle\n",
    "import argparse\n",
    "from transformers import AutoTokenizer, DistilBertForSequenceClassification, DistilBertModel\n",
    "\n",
    "#### Just some code to print debug information to stdout\n",
    "logging.basicConfig(format='%(asctime)s - %(message)s',\n",
    "                    datefmt='%Y-%m-%d %H:%M:%S',\n",
    "                    level=logging.INFO,\n",
    "                    handlers=[LoggingHandler()])\n",
    "#### /print debug information to stdout\n",
    "\n",
    "\n",
    "# parser = argparse.ArgumentParser()\n",
    "# parser.add_argument(\"--train_batch_size\", default=64, type=int)\n",
    "# parser.add_argument(\"--max_seq_length\", default=512, type=int)\n",
    "# parser.add_argument(\"--model_name\", required=True)\n",
    "# parser.add_argument(\"--max_passages\", default=0, type=int)\n",
    "# parser.add_argument(\"--epochs\", default=10, type=int)\n",
    "# parser.add_argument(\"--pooling\", default=\"mean\")\n",
    "# parser.add_argument(\"--negs_to_use\", default=None, help=\"From which systems should negatives be used? Multiple systems seperated by comma. None = all\")\n",
    "# parser.add_argument(\"--warmup_steps\", default=1000, type=int)\n",
    "# parser.add_argument(\"--lr\", default=2e-5, type=float)\n",
    "# parser.add_argument(\"--num_negs_per_system\", default=5, type=int)\n",
    "# parser.add_argument(\"--use_pre_trained_model\", default=False, action=\"store_true\")\n",
    "# parser.add_argument(\"--use_all_queries\", default=False, action=\"store_true\")\n",
    "# parser.add_argument(\"--ce_score_margin\", default=3.0, type=float)\n",
    "# args = parser.parse_args()\n",
    "\n",
    "# print(args)\n",
    "\n",
    "# The  model we want to fine-tune\n",
    "# model_name = args.model_name\n",
    "\n",
    "# train_batch_size = args.train_batch_size           #Increasing the train batch size improves the model performance, but requires more GPU memory\n",
    "# max_seq_length = args.max_seq_length            #Max length for passages. Increasing it, requires more GPU memory\n",
    "# ce_score_margin = args.ce_score_margin             #Margin for the CrossEncoder score between negative and positive passages\n",
    "# num_negs_per_system = args.num_negs_per_system         # We used different systems to mine hard negatives. Number of hard negatives to add from each system\n",
    "# num_epochs = args.epochs                 # Number of epochs we want to train\n",
    "model_name = \"distilbert-base-uncased\"\n",
    "# Distill_classical_model = DistilBertModel.from_pretrained(\"distilbert-base-uncased\")\n",
    "# state_dict = Distill_classical_model.state_dict()\n",
    "\n",
    "train_batch_size = 1           #Increasing the train batch size improves the model performance, but requires more GPU memory\n",
    "max_seq_length = 512            #Max length for passages. Increasing it, requires more GPU memory\n",
    "# ce_score_margin = args.ce_score_margin             #Margin for the CrossEncoder score between negative and positive passages\n",
    "# num_negs_per_system = args.num_negs_per_system         # We used different systems to mine hard negatives. Number of hard negatives to add from each system\n",
    "num_epochs = 10                 # Number of epochs we want to train\n",
    "\n",
    "use_pre_trained_model = True\n",
    "# Load our embedding model\n",
    "if use_pre_trained_model:\n",
    "    logging.info(\"use pretrained SBERT model\")\n",
    "    model = SentenceTransformer(model_name)\n",
    "    model = model.load_state_dict(state_dict, )\n",
    "    model.max_seq_length = max_seq_length\n",
    "else:\n",
    "    logging.info(\"Create new SBERT model\")\n",
    "    word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)\n",
    "    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), args.pooling)\n",
    "    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])\n",
    "\n",
    "model_save_path = 'output/train_bi-encoder-mnrl-{}-{}-pretrain-classical'.format(model_name.replace(\"/\", \"-\"),  datetime.now().strftime(\"%Y-%m-%d_%H-%M-%S\"))\n",
    "\n",
    "\n",
    "### Now we read the MS Marco dataset\n",
    "# data_folder = 'msmarco-data'\n",
    "\n",
    "# For training the SentenceTransformer model, we need a dataset, a dataloader, and a loss used for training.\n",
    "# train_dataset = MSMARCODataset(train_queries, corpus=corpus)\n",
    "train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=train_batch_size)\n",
    "train_loss = losses.MultipleNegativesRankingLoss(model=model)\n",
    "\n",
    "# Train the model\n",
    "model.fit(train_objectives=[(train_dataloader, train_loss)],\n",
    "          epochs=num_epochs,\n",
    "          warmup_steps=1000,\n",
    "          use_amp=True,\n",
    "          checkpoint_path=model_save_path,\n",
    "          checkpoint_save_steps=len(train_dataloader),\n",
    "          optimizer_params = {'lr': 2e-5},\n",
    "          )\n",
    "\n",
    "# Save the model\n",
    "# model.save(model_save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae06dd75-8d12-4fe3-ad28-0626d840f013",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd7f7a30-d9af-4c4a-a507-d3958dda0528",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(range(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8b1db4e-c847-4ba6-8052-8916f561d2c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"dev.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c608d69a-d5f4-49b5-9e67-f7ef9c9d58a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"dev.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"random_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e525d01c-eb3f-4c03-bf1c-af7427f18029",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(context,pos,neg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f62e074-426d-4173-b121-ddeff59ec2cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-bert-base-uncased-2023-08-29_15-22-44/86808\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cb31a6d-d24b-4b2a-a254-507fb4b1267d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-random-bert-base-uncased-2023-08-29_15-26-52\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b86a24-2d9f-4179-946b-f58475aee420",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pos+adv训练，pos+random测试\n",
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"random_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-bert-base-uncased-2023-08-29_15-22-44/86808\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bf1f8e7-3f2c-4848-b6b6-84d818ce1745",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-distilbert-base-uncased-2023-08-23_21-32-34/72340\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ca285a8-cb45-441a-9505-3ed64681cc2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"random_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-distilbert-base-uncased-2023-08-23_21-32-34/72340\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8af3235-8544-4975-b746-afdd0869dfaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-random-distilbert-base-uncased-2023-08-24_10-25-55/72340\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "634c2675-d342-40db-b050-bf71f22f8d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-sentence-t5-base-2023-08-30_11-35-59\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95be0dda-3eb1-488b-bafa-14ebd3085b89",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            context_list.append(context)\n",
    "            pos_list.append(pos)\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            neg_list.append(neg)\n",
    "            dev_examples.append(InputExample(texts=[context, pos, neg]))\n",
    "evaluator = evaluation.TripletEvaluator(context_list, pos_list, neg_list, 0)\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-random-sentence-t5-base-2023-08-30_11-36-59\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68269aa-8825-485c-86cf-4959dfb7eef0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "line = dev_lines[5]\n",
    "y_list = []\n",
    "\n",
    "# print(line)\n",
    "sample=json.loads(line)\n",
    "ss = \"[CLS] \"\n",
    "for c in sample[\"context\"]:\n",
    "    ss = ss + c + \" [SEP] \"\n",
    "context = ss.strip()\n",
    "y_list.append(0)\n",
    "for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "    pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "    pos_list.append(pos)\n",
    "    y_list.append(1)\n",
    "for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "    neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "    neg_list.append(neg)\n",
    "    y_list.append(2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fa1e9b3-96fa-4ade-83c5-96c8efdf416b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"context:\")\n",
    "print(context)\n",
    "print(\"\\n\")\n",
    "\n",
    "print(\"pos:\")\n",
    "for i, pos in enumerate(pos_list):\n",
    "    print(i, pos)\n",
    "print(\"\\n\")\n",
    "\n",
    "print(\"neg:\")\n",
    "for i, neg in enumerate(neg_list):\n",
    "    print(i, neg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee682dd1-4d05-48ba-8625-ae750f591075",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-distilbert-base-uncased-2023-08-23_21-32-34/72340\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05db2db5-ed09-4700-b52d-3e334671521e",
   "metadata": {},
   "outputs": [],
   "source": [
    "context_embedding = model.encode(context)\n",
    "print(context_embedding.shape)\n",
    "x_list = []\n",
    "x_list.append(context_embedding)\n",
    "for pos in pos_list:\n",
    "    pos_embedding = model.encode(pos)\n",
    "    x_list.append(pos_embedding)\n",
    "for neg in neg_list:\n",
    "    neg_embedding = model.encode(neg)\n",
    "    x_list.append(neg_embedding)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e88e82b1-3c78-4fc2-b8e6-21009ff9a2f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(x_list), len(y_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15ac5cc3-162f-4ba2-846f-8c8557a465fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "x = np.array(x_list)\n",
    "y = np.array(y_list)\n",
    "print(len(x), len(y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e0905c0-52c5-4aba-b071-65e9094c5165",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from openTSNE import TSNE\n",
    "from openTSNE.affinity import PerplexityBasedNN\n",
    "from openTSNE import initialization\n",
    "# from openTSNE.callbacks import ErrorLogger\n",
    "import utils\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "802b892f-dfa8-4e1c-9370-35411c03f147",
   "metadata": {},
   "outputs": [],
   "source": [
    "tsne = TSNE(\n",
    "    perplexity=50,\n",
    "    n_iter=200,\n",
    "    metric=\"cosine\",\n",
    "    n_jobs=8,\n",
    "    random_state=0,\n",
    ")\n",
    "embedding = tsne.fit(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8137023-e294-4be3-af69-e522806cc9a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = embedding\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "\n",
    "classes = np.unique(y)\n",
    "\n",
    "default_colors = matplotlib.rcParams[\"axes.prop_cycle\"]\n",
    "colors = {k: v[\"color\"] for k, v in zip(classes, default_colors())}\n",
    "\n",
    "point_colors = list(map(colors.get, y))\n",
    "\n",
    "ax.scatter(x[:, 0], x[:, 1], c=point_colors, marker=\"v\", s=50)\n",
    "\n",
    "legend_handles = [\n",
    "    matplotlib.lines.Line2D(\n",
    "        [],\n",
    "        [],\n",
    "        marker=\"s\",\n",
    "        color=\"w\",\n",
    "        markerfacecolor=colors[yi],\n",
    "        ms=10,\n",
    "        alpha=1,\n",
    "        linewidth=0,\n",
    "        markeredgecolor=\"k\",\n",
    "#         label=yi\n",
    "    )\n",
    "    for yi in classes\n",
    "]\n",
    "legend_kwargs_ = dict(labels=[\"context\", \"pos\", \"neg\"],bbox_to_anchor=(0.32, 1.18), frameon=False, fontsize=10)\n",
    "ax.legend(handles=legend_handles, **legend_kwargs_)\n",
    "ax.set_title(\"example\")\n",
    "fig.savefig('dd_similar.png',bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "161b162a-ce48-4647-bbfe-5fe1a8a99081",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics.pairwise import paired_cosine_distances, paired_euclidean_distances, paired_manhattan_distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f70eff66-2e0c-49f6-ac9a-65f500553049",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    dev_lines = f.readlines()\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-distilbert-base-uncased-2023-08-23_21-32-34/72340\")\n",
    "dev_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in dev_lines:\n",
    "    # y_list = []\n",
    "    \n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    # y_list.append(0)\n",
    "    pos_dis_list = []\n",
    "    context_embedding = model.encode(context, convert_to_numpy=True)\n",
    "    context_embedding = context_embedding.reshape(1,-1)\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        pos_embedding = model.encode(pos, convert_to_numpy=True)\n",
    "        pos_embedding = pos_embedding.reshape(1,-1)\n",
    "        pos_cos_distance = paired_cosine_distances(context_embedding, pos_embedding)\n",
    "        pos_dis_list.append(pos_cos_distance)\n",
    "    \n",
    "    neg_dis_list = []\n",
    "    for j, n in enumerate(sample[\"random_negative_responses\"]):\n",
    "        neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "        neg_embedding = model.encode(neg, convert_to_numpy=True)\n",
    "        neg_embedding = neg_embedding.reshape(1,-1)\n",
    "        neg_cos_distance = paired_cosine_distances(context_embedding, neg_embedding)\n",
    "        neg_dis_list.append(neg_cos_distance)\n",
    "    pos_list.append(pos_dis_list)\n",
    "    neg_list.append(neg_dis_list)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9a803a4-0752-4e54-a1fd-cbe8ea625a2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_pos_mean_list = []\n",
    "total_neg_mean_list = []\n",
    "for i, pos_dis in enumerate(pos_list):\n",
    "    dis_sum = sum(pos_dis)\n",
    "    total_pos_mean_list.append(dis_sum/len(pos_dis))\n",
    "\n",
    "for i, neg_dis in enumerate(neg_list):\n",
    "    dis_sum = sum(neg_dis)\n",
    "    total_neg_mean_list.append(dis_sum/len(neg_dis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08230238-2f33-40fc-9e5f-e2fe6af1d3d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(total_pos_mean_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb55f25e-60c6-4f8d-90b9-1523a282348f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(total_neg_mean_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ad3f1f1-2612-42f1-ad17-82a35a92d1bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(total_neg_mean_list[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "894b258c-b59b-49ec-9de2-b850993d1980",
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "for dis in total_pos_mean_list:\n",
    "    if dis[0] > 0.9:\n",
    "        count = count + 1\n",
    "print(count, len(total_pos_mean_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f8c13b7-a2f3-43b0-aeeb-b3b2a1db967d",
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "for dis in total_neg_mean_list:\n",
    "    if dis[0] < 0.9:\n",
    "        count = count + 1\n",
    "print(count, len(total_neg_mean_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f784cea1-7991-4e67-bf40-b76e1ed49f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "n = 0\n",
    "for i, pos_dis in enumerate(pos_list):\n",
    "    for dis in pos_dis:\n",
    "        n = n + 1\n",
    "        if dis[0] > 1.2:\n",
    "            count = count + 1\n",
    "\n",
    "for i, neg_dis in enumerate(neg_list):\n",
    "    for dis in neg_dis:\n",
    "        n = n + 1\n",
    "        if dis[0] < 0.3:\n",
    "            count = count + 1\n",
    "print(count, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e0a8a1c-4458-4902-930e-bb949f62d936",
   "metadata": {},
   "outputs": [],
   "source": [
    "5/11420"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2ef008e-27c7-4a3c-8a25-1469819be23d",
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "n = 0\n",
    "for i, pos_dis in enumerate(pos_list):\n",
    "    for neg_dis in neg_list[i]:\n",
    "        if neg[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c984910b-c2e1-48eb-ad10-78bd25cd8daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "from typing import Iterable, Dict\n",
    "import torch.nn.functional as F\n",
    "from torch import nn, Tensor\n",
    "from sentence_transformers.SentenceTransformer import SentenceTransformer\n",
    "from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, hidden_size, num_labels):\n",
    "        super(MLP, self).__init__()\n",
    "        self.classifier = nn.Linear(hidden_size, num_labels)\n",
    "        self.num_labels = num_labels\n",
    "        \n",
    "    def forward(self, hidden_states, labels):\n",
    "        logits = self.classifier(hidden_states)\n",
    "\n",
    "        loss_fct = CrossEntropyLoss()\n",
    "        loss = loss_fct(logits, labels)\n",
    "        return logits, loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a44e205-6829-493f-b8be-322ec5e60cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"train.json\", \"r\") as f:\n",
    "    train_lines = f.readlines()\n",
    "train_classification_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in train_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        train_classification_examples.append((context, pos, 1))\n",
    "    for j, n in enumerate(sample[\"random_negative_responses\"]):\n",
    "        neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "        train_classification_examples.append((context, neg, 0))\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-distilbert-base-uncased-2023-08-23_21-32-34/72340\")\n",
    "# print(train_classification_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecc9cf08-fded-4e88-b790-4a1aa45aa787",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"test.json\", \"r\") as f:\n",
    "    test_lines = f.readlines()\n",
    "test_classification_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in test_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        test_classification_examples.append((context, pos, 1))\n",
    "    for j, n in enumerate(sample[\"random_negative_responses\"]):\n",
    "        neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "        test_classification_examples.append((context, neg, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7bd4a32-0b1c-4d33-9597-c23275290d38",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_classification_dataloader = DataLoader(train_classification_examples, shuffle=True, batch_size=32768)\n",
    "test_classification_dataloader = DataLoader(test_classification_examples, shuffle=True, batch_size=32768)\n",
    "num_train_data = len(train_classification_examples)\n",
    "num_test_data = len(test_classification_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c39c958f-7adb-4fd3-82d5-d967bf66e555",
   "metadata": {},
   "outputs": [],
   "source": [
    "for data in train_classification_dataloader:\n",
    "    print(data)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d67298ca-25bb-4cb9-876a-be1342549669",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_classification_examples[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31928e25-1e03-4254-aa8a-ca4a9fa985a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = train_classification_examples[0]\n",
    "with torch.no_grad():\n",
    "    context_embedding = torch.tensor(model.encode(sample[0]))\n",
    "    response_embedding = torch.tensor(model.encode(sample[1]))\n",
    "    print(context_embedding.shape)\n",
    "    hidden_state = torch.cat([context_embedding, response_embedding])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd68b127-7c80-4605-9f66-e311e7b0317a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(hidden_state.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56f2ecae-b043-4928-aab5-6227efa1f2c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "classification_model = MLP(1536, 2).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d368257-4791-4a4b-ab61-8ee74a71952c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "num_epoch = 30\n",
    "epoch = 0\n",
    "load_checkpoint_path = \"./random_test/\"\n",
    "if not os.path.exists(load_checkpoint_path):\n",
    "   os.mkdir(load_checkpoint_path)\n",
    "optimizer = torch.optim.Adam(classification_model.parameters(), lr=1e-5)\n",
    "for epoch in range(num_epoch):\n",
    "    accu = 0\n",
    "    for train_data in train_classification_dataloader:\n",
    "        with torch.no_grad():\n",
    "            context_embedding = torch.tensor(model.encode(train_data[0])).to(\"cuda\")\n",
    "            response_embedding = torch.tensor(model.encode(train_data[1])).to(\"cuda\")\n",
    "            # print(context_embedding.shape)\n",
    "            hidden_state = torch.cat([context_embedding, response_embedding], dim=1)\n",
    "            # print(hidden_state.shape)\n",
    "        \n",
    "        outputs = classification_model(hidden_state, train_data[2].to(\"cuda\"))\n",
    "        # print(data[2])\n",
    "        loss = outputs[1]\n",
    "        logits = outputs[0]\n",
    "        # print(outputs)\n",
    "        # print(loss)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    total_num_right = 0\n",
    "    for test_data in test_classification_dataloader:\n",
    "        with torch.no_grad():\n",
    "            labels = test_data[2].to(\"cuda\")\n",
    "            context_embedding = torch.tensor(model.encode(test_data[0])).to(\"cuda\")\n",
    "            response_embedding = torch.tensor(model.encode(test_data[1])).to(\"cuda\")\n",
    "            # print(context_embedding.shape)\n",
    "            hidden_state = torch.cat([context_embedding, response_embedding], dim=1)\n",
    "            # print(hidden_state.shape)\n",
    "            outputs = classification_model(hidden_state, labels)\n",
    "            logits = outputs[0]\n",
    "            score = logits.argmax(dim=-1)\n",
    "            num_right = ((score == labels).float()).sum()\n",
    "            total_num_right += num_right\n",
    "    right_rate = total_num_right/num_test_data\n",
    "    print(right_rate)\n",
    "    if right_rate > accu:\n",
    "        accu = right_rate\n",
    "        state = {\n",
    "            \"epoch\": epoch,\n",
    "            \"model_state_dict\": model.state_dict(),\n",
    "            \"best_accu\": accu,\n",
    "        }\n",
    "        torch.save(state, load_checkpoint_path + \"checkpoint.bin\")\n",
    "    # print(right_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c97bc3c-d1d2-4b62-a42f-b706218f468a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_examples = []\n",
    "for line in lines[:10]:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            train_examples.append([context, pos, neg])\n",
    "train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9172cc2-458b-4404-89eb-5cc7b932a012",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, data in enumerate(train_dataloader):\n",
    "    pos_embed = "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e09a1af-0f9d-42a3-9d82-6795f9fa1ed2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer, LoggingHandler, losses, InputExample\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "model = SentenceTransformer(\"output/train_bi-encoder-mnrl-bert-base-uncased-2023-08-29_15-22-44/86808\")\n",
    "train_examples = [\n",
    "    InputExample(texts=['This is a positive pair', 'Where the distance will be minimized'], label=1),\n",
    "    InputExample(texts=['This is a negative pair', 'Their distance will be increased'], label=0)]\n",
    "\n",
    "train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=2)\n",
    "train_loss = losses.ContrastiveLoss(model=model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3326add-223b-4ed8-b632-f81d4f427132",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_dataloader)\n",
    "data_iterators = iter(train_dataloader)\n",
    "data = next(data_iterators)\n",
    "\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "808c176a-0670-4dec-952f-9d432164952a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open (\"train.json\", \"r\") as f:\n",
    "    lines = f.readlines()\n",
    "\n",
    "train_examples = []\n",
    "for line in lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "            neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "            train_examples.append((context, pos, neg))\n",
    "model_encoder = SentenceTransformer(\"output/train_bi-encoder-mnrl-bert-base-uncased-2023-08-29_15-22-44/86808\")\n",
    "# print(train_classification_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d422aa41-e014-4b5a-be68-c86062ebb81e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SiameseDistanceMetric(Enum):\n",
    "    \"\"\"\n",
    "    The metric for the contrastive loss\n",
    "    \"\"\"\n",
    "    EUCLIDEAN = lambda x, y: F.pairwise_distance(x, y, p=2)\n",
    "    MANHATTAN = lambda x, y: F.pairwise_distance(x, y, p=1)\n",
    "    COSINE_DISTANCE = lambda x, y: 1-F.cosine_similarity(x, y, dim=0)\n",
    "\n",
    "\n",
    "class test_ContrastiveLoss(nn.Module):\n",
    "    \"\"\"\n",
    "    Contrastive loss. Expects as input two texts and a label of either 0 or 1. If the label == 1, then the distance between the\n",
    "    two embeddings is reduced. If the label == 0, then the distance between the embeddings is increased.\n",
    "\n",
    "    Further information: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf\n",
    "\n",
    "    :param model: SentenceTransformer model\n",
    "    :param distance_metric: Function that returns a distance between two embeddings. The class SiameseDistanceMetric contains pre-defined metrices that can be used\n",
    "    :param margin: Negative samples (label == 0) should have a distance of at least the margin value.\n",
    "    :param size_average: Average by the size of the mini-batch.\n",
    "\n",
    "    Example::\n",
    "\n",
    "        from sentence_transformers import SentenceTransformer, LoggingHandler, losses, InputExample\n",
    "        from torch.utils.data import DataLoader\n",
    "\n",
    "        model = SentenceTransformer('all-MiniLM-L6-v2')\n",
    "        train_examples = [\n",
    "            InputExample(texts=['This is a positive pair', 'Where the distance will be minimized'], label=1),\n",
    "            InputExample(texts=['This is a negative pair', 'Their distance will be increased'], label=0)]\n",
    "\n",
    "        train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=2)\n",
    "        train_loss = losses.ContrastiveLoss(model=model)\n",
    "\n",
    "        model.fit([(train_dataloader, train_loss)], show_progress_bar=True)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, distance_metric=SiameseDistanceMetric.COSINE_DISTANCE, margin: float = 0.5, size_average:bool = True):\n",
    "        super(test_ContrastiveLoss, self).__init__()\n",
    "        self.distance_metric = distance_metric\n",
    "        self.margin = margin\n",
    "        self.size_average = size_average\n",
    "\n",
    "    def get_config_dict(self):\n",
    "        distance_metric_name = self.distance_metric.__name__\n",
    "        for name, value in vars(SiameseDistanceMetric).items():\n",
    "            if value == self.distance_metric:\n",
    "                distance_metric_name = \"SiameseDistanceMetric.{}\".format(name)\n",
    "                break\n",
    "\n",
    "        return {'distance_metric': distance_metric_name, 'margin': self.margin, 'size_average': self.size_average}\n",
    "\n",
    "    def forward(self, reps, labels: Tensor):\n",
    "        # reps = [self.model(sentence_feature)['sentence_embedding'] for sentence_feature in sentence_features]\n",
    "        # assert len(reps) == 2\n",
    "        rep_anchor, rep_other = reps.chunk(2,-1)\n",
    "        print(rep_anchor.shape,rep_other.shape)\n",
    "        distances = self.distance_metric(rep_anchor, rep_other)\n",
    "        losses = 0.5 * (labels.float() * distances.pow(2) + (1 - labels).float() * F.relu(self.margin - distances).pow(2))\n",
    "        return losses.mean() if self.size_average else losses.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3f58f5e-fe86-4292-8201-88e4787f802d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = train_examples[0]\n",
    "print(sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "412d2cc2-ac73-4580-aceb-8682e6d2579e",
   "metadata": {},
   "outputs": [],
   "source": [
    "context_embedding = torch.tensor(model_encoder.encode(sample[0]))\n",
    "pos_response_embedding = torch.tensor(model_encoder.encode(sample[1]))\n",
    "neg_response_embedding = torch.tensor(model_encoder.encode(sample[2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50403942-1dcf-4df7-953e-f0a9aeba10e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(type(model_encoder.encode(sample[0])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d44548a-e48c-4b1c-8aeb-945eb577e781",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss = test_ContrastiveLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c726bf4f-958f-45df-be91-4a2afd6bd3b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "loss = test_loss(pos_response_embedding, labels=torch.tensor(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6973b93b-10fc-4729-8b31-003a9faf283f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33396df1-f132-465f-bb66-80111d00a837",
   "metadata": {},
   "outputs": [],
   "source": [
    "from "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c7da553-9c98-4c73-9a69-1af788c7402b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(robust_pos_response_feature.shape, pos_response_embedding.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b69aef54-ddba-409a-b79a-31130ef1c836",
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "from typing import Iterable, Dict\n",
    "import torch.nn.functional as F\n",
    "from torch import nn, Tensor\n",
    "from sentence_transformers.SentenceTransformer import SentenceTransformer\n",
    "\n",
    "\n",
    "class SiameseDistanceMetric(Enum):\n",
    "    \"\"\"\n",
    "    The metric for the contrastive loss\n",
    "    \"\"\"\n",
    "    EUCLIDEAN = lambda x, y: F.pairwise_distance(x, y, p=2)\n",
    "    MANHATTAN = lambda x, y: F.pairwise_distance(x, y, p=1)\n",
    "    COSINE_DISTANCE = lambda x, y: 1-F.cosine_similarity(x, y)\n",
    "\n",
    "\n",
    "class DisContrastiveLoss(nn.Module):\n",
    "    \"\"\"\n",
    "    Contrastive loss. Expects as input two texts and a label of either 0 or 1. If the label == 1, then the distance between the\n",
    "    two embeddings is reduced. If the label == 0, then the distance between the embeddings is increased.\n",
    "\n",
    "    Further information: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf\n",
    "\n",
    "    :param model: SentenceTransformer model\n",
    "    :param distance_metric: Function that returns a distance between two embeddings. The class SiameseDistanceMetric contains pre-defined metrices that can be used\n",
    "    :param margin: Negative samples (label == 0) should have a distance of at least the margin value.\n",
    "    :param size_average: Average by the size of the mini-batch.\n",
    "\n",
    "    Example::\n",
    "\n",
    "        from sentence_transformers import SentenceTransformer, LoggingHandler, losses, InputExample\n",
    "        from torch.utils.data import DataLoader\n",
    "\n",
    "        model = SentenceTransformer('all-MiniLM-L6-v2')\n",
    "        train_examples = [\n",
    "            InputExample(texts=['This is a positive pair', 'Where the distance will be minimized'], label=1),\n",
    "            InputExample(texts=['This is a negative pair', 'Their distance will be increased'], label=0)]\n",
    "\n",
    "        train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=2)\n",
    "        train_loss = losses.ContrastiveLoss(model=model)\n",
    "\n",
    "        model.fit([(train_dataloader, train_loss)], show_progress_bar=True)\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(self, distance_metric=SiameseDistanceMetric.COSINE_DISTANCE, margin: float = 0.5, size_average:bool = True):\n",
    "        super(DisContrastiveLoss, self).__init__()\n",
    "        self.distance_metric = distance_metric\n",
    "        self.margin = margin\n",
    "        # self.model = model\n",
    "        self.size_average = size_average\n",
    "\n",
    "    def get_config_dict(self):\n",
    "        distance_metric_name = self.distance_metric.__name__\n",
    "        for name, value in vars(SiameseDistanceMetric).items():\n",
    "            if value == self.distance_metric:\n",
    "                distance_metric_name = \"SiameseDistanceMetric.{}\".format(name)\n",
    "                break\n",
    "\n",
    "        return {'distance_metric': distance_metric_name, 'margin': self.margin, 'size_average': self.size_average}\n",
    "\n",
    "    def forward(self, pos_response_feature, neg_response_feature):\n",
    "        # reps = [self.model(sentence_feature)['sentence_embedding'] for sentence_feature in sentence_features]\n",
    "        # assert len(reps) == 2\n",
    "        robust_pos_response_feature, non_robust_pos_response_feature = pos_response_feature.chunk(2,-1)\n",
    "        robust_neg_response_feature, non_robust_neg_response_feature = neg_response_feature.chunk(2,-1)\n",
    "        # 使同 一 feature里的robust与non robust feature远离\n",
    "        pos_feature_distances = self.distance_metric(robust_pos_response_feature, non_robust_pos_response_feature)\n",
    "        inside_pos_losses = 0.5 * (0.float() * pos_feature_distances.pow(2) + (1 - 0).float() * F.relu(self.margin - pos_feature_distances).pow(2))\n",
    "        \n",
    "        neg_feature_distances = self.distance_metric(robust_neg_response_feature, non_robust_neg_response_feature)\n",
    "        inside_neg_losses = 0.5 * (0.float() * neg_feature_distances.pow(2) + (1 - 0).float() * F.relu(self.margin - neg_feature_distances).pow(2))\n",
    "        \n",
    "        #使不同feature的robust接近\n",
    "        robust_distances = self.distance_metric(robust_neg_response_feature, robust_neg_response_feature)\n",
    "        \n",
    "        robust_losses = 0.5 * (labels.float() * robust_distances.pow(2) + (1 - labels).float() * F.relu(self.margin - robust_distances).pow(2))\n",
    "        non_robust_losses = 0.5 * (labels.float() * non_robust_distances.pow(2) + (1 - labels).float() * F.relu(self.margin - non_robust_distances).pow(2))\n",
    "        return losses.mean() if self.size_average else losses.sum()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7c0babb-6bb9-426c-84ff-2de53d046c80",
   "metadata": {},
   "outputs": [],
   "source": [
    "constrastive_loss = DisContrastiveLoss()\n",
    "loss = constrastive_loss(pos_response_embedding, neg_response_embedding)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c011ffc7-210c-4adb-af2f-e53d4c828b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "azime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac6b21d7-1970-4c43-a586-69ea30db6de4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!conda list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8e0db7c-2670-4bc9-b825-c71f8b94bf4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "a = torch.ones(32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2230bd2f-0846-4250-96c3-45b01530f41e",
   "metadata": {},
   "outputs": [],
   "source": [
    "2*a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67061ccd-e3e8-44ce-9964-295cdc172bfc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea2bfad3-7a49-4873-9ffd-6b37b352941a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# no contrastive learning accuracy\n",
    "with open (\"train.json\", \"r\") as f:\n",
    "    train_lines = f.readlines()\n",
    "train_classification_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in train_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        train_classification_examples.append((context, pos, 1))\n",
    "    for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "        neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "        train_classification_examples.append((context, neg, 0))\n",
    "        \n",
    "with open (\"test.json\", \"r\") as f:\n",
    "    test_lines = f.readlines()\n",
    "test_classification_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in test_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        test_classification_examples.append((context, pos, 1))\n",
    "    for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "        neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "        test_classification_examples.append((context, neg, 0))\n",
    "train_classification_dataloader = DataLoader(train_classification_examples, shuffle=True, batch_size=32)\n",
    "test_classification_dataloader = DataLoader(test_classification_examples, shuffle=True, batch_size=32)\n",
    "num_train_data = len(train_classification_examples)\n",
    "num_test_data = len(test_classification_examples)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa2aa1c-8b23-48d4-aad9-767c7fc789f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for data in train_classification_dataloader:\n",
    "    print(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f86497cd-7d45-4ced-bcbc-0f8bf07a64fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, DistilBertForSequenceClassification, DistilBertModel\n",
    "from enum import Enum\n",
    "from typing import Iterable, Dict\n",
    "import torch.nn.functional as F\n",
    "from torch import nn, Tensor\n",
    "from sentence_transformers.SentenceTransformer import SentenceTransformer\n",
    "from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss\n",
    "class classifier_no_contrastive(nn.Module):\n",
    "    def __init__(self, hidden_size, num_labels):\n",
    "        super(classifier_no_contrastive, self).__init__()\n",
    "        self.encoder = DistilBertModel.from_pretrained(\"distilbert-base-uncased\")\n",
    "        self.classifier = nn.Linear(hidden_size, num_labels)\n",
    "        self.num_labels = num_labels\n",
    "        \n",
    "    def forward(self, data):\n",
    "        contexts = tokenizer(data[0], return_tensors=\"pt\", padding=True, truncation=True)\n",
    "        responses = tokenizer(data[1], return_tensors=\"pt\", padding=True)\n",
    "        context_embeddings = self.encoder(**contexts)[0][:, 0]\n",
    "        responses_embeddings = self.encoder(**responses)[0][:, 0]\n",
    "        hidden_states = torch.cat([context_embeddings, responses_embeddings],dim=1)\n",
    "        labels = data[2]\n",
    "        # hidden_states = self.encoder(inputs)\n",
    "        logits = self.classifier(hidden_states)\n",
    "\n",
    "        loss_fct = CrossEntropyLoss()\n",
    "        loss = loss_fct(logits, labels)\n",
    "        return logits, loss\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# predicted_class_id = logits.argmax().item()\n",
    "\n",
    "# # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`\n",
    "# num_labels = len(model.config.id2label)\n",
    "# model = DistilBertForSequenceClassification.from_pretrained(\"distilbert-base-uncased\", num_labels=num_labels)\n",
    "\n",
    "# labels = torch.tensor([1])\n",
    "# loss = model(**inputs, labels=labels).loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "058b13e9-05c5-4346-a2ea-c51c0dd3e39e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, DistilBertForSequenceClassification, DistilBertModel\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "model = DistilBertModel.from_pretrained(\"distilbert-base-uncased\")\n",
    "for data in train_classification_dataloader:\n",
    "    contexts = tokenizer(data[0], return_tensors=\"pt\", padding=True, truncation=True)\n",
    "    responses = tokenizer(data[1], return_tensors=\"pt\", padding=True)\n",
    "    context_embeddings = model(**contexts)[0][:, 0]\n",
    "    responses_embeddings = model(**responses)[0][:, 0]\n",
    "    hidden_states = torch.cat([context_embeddings, responses_embeddings],dim=1)\n",
    "    print(hidden_states.shape)\n",
    "# with torch.no_grad():\n",
    "#     logits = model(**inputs).logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df26772d-bdae-4385-881a-c289a0b1ceae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "num_epoch = 30\n",
    "epoch = 0\n",
    "load_checkpoint_path = \"./pretrain_basic_train_basic/\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "classification_model_1 = classifier_no_contrastive(1536, 2)\n",
    "if not os.path.exists(load_checkpoint_path):\n",
    "   os.mkdir(load_checkpoint_path)\n",
    "optimizer = torch.optim.Adam(classification_model.parameters(), lr=1e-5)\n",
    "for epoch in range(num_epoch):\n",
    "    accu = 0\n",
    "    for train_data in train_classification_dataloader:\n",
    "        # with torch.no_grad():\n",
    "        #     context_embedding = torch.tensor(model.encode(train_data[0]))\n",
    "        #     response_embedding = torch.tensor(model.encode(train_data[1]))\n",
    "        #     # print(context_embedding.shape)\n",
    "        #     hidden_state = torch.cat([context_embedding, response_embedding], dim=1)\n",
    "        #     # print(hidden_state.shape)\n",
    "\n",
    "        outputs = classification_model_1(train_data)\n",
    "        # print(data[2])\n",
    "        loss = outputs[1]\n",
    "        logits = outputs[0]\n",
    "        # print(outputs)\n",
    "        # print(loss)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    total_num_right = 0\n",
    "    for test_data in test_classification_dataloader:\n",
    "        with torch.no_grad():\n",
    "            # labels = test_data[2]\n",
    "            # context_embedding = torch.tensor(model.encode(test_data[0]))\n",
    "            # response_embedding = torch.tensor(model.encode(test_data[1]))\n",
    "            # # print(context_embedding.shape)\n",
    "            # hidden_state = torch.cat([context_embedding, response_embedding], dim=1)\n",
    "            # print(hidden_state.shape)\n",
    "            outputs = classification_model_1(test_data)\n",
    "            logits = outputs[0]\n",
    "            score = logits.argmax(dim=-1)\n",
    "            num_right = ((score == labels).float()).sum()\n",
    "            total_num_right += num_right\n",
    "    right_rate = total_num_right/num_test_data\n",
    "    print(right_rate)\n",
    "    if right_rate > accu:\n",
    "        accu = right_rate\n",
    "        state = {\n",
    "            \"epoch\": epoch,\n",
    "            \"model_state_dict\": model.state_dict(),\n",
    "            \"best_accu\": accu,\n",
    "        }\n",
    "        torch.save(state, load_checkpoint_path + \"checkpoint.bin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b198b3d4-85c3-4c2a-a67a-46e25a9a9a5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "with open (\"test.json\", \"r\") as f:\n",
    "    test_lines = f.readlines()\n",
    "test_classification_examples = []\n",
    "context_list = []\n",
    "pos_list = []\n",
    "neg_list = []\n",
    "for line in test_lines:\n",
    "    # print(line)\n",
    "    sample=json.loads(line)\n",
    "    ss = \"[CLS] \"\n",
    "    for c in sample[\"context\"]:\n",
    "        ss = ss + c + \" [SEP] \"\n",
    "    context = ss.strip()\n",
    "    for i, p in enumerate(sample[\"positive_responses\"]):\n",
    "        pos = \"[CLS] \" + p + \" [SEP]\"\n",
    "        test_classification_examples.append((context, pos, 1))\n",
    "    for j, n in enumerate(sample[\"adversarial_negative_responses\"]):\n",
    "        neg = \"[CLS] \" + n + \" [SEP]\"\n",
    "        test_classification_examples.append((context, neg, 0))\n",
    "random.shuffle(test_classification_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40f5cd16-37da-47a7-853d-d1bbf4ab1d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "json_samples = json.dumps(test_classification_examples, ensure_ascii=False, indent=2)\n",
    "with open(\"testset.json\", 'w', encoding='utf-8') as f:\n",
    "    f.write(json_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c591c8d-c8a2-474e-b91e-aca12a512e39",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "disentangle",
   "language": "python",
   "name": "disentangle"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}