sbert_relevance_scores.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from rank_bm25 import BM25Okapi\n",
    "import os\n",
    "import numpy\n",
    "from scripts.utils import read_gold_data\n",
    "import numpy as np\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path_name = \"data-release\" # set data path name!\n",
    "\n",
    "data = read_gold_data(data_path_name)\n",
    "corpus = data[\"corpus\"]\n",
    "baseline_queries_train = data[\"baseline\"][\"train\"]\n",
    "baseline_queries_dev = data[\"baseline\"][\"dev\"]\n",
    "perspective_queries_train = data[\"perspective\"][\"train\"]\n",
    "perspective_queries_dev = data[\"perspective\"][\"dev\"]\n",
    "\n",
    "# uncomment when using test set\n",
    "# baseline_queries_test = data[\"baseline\"][\"test\"]\n",
    "# perspective_queries_test = data[\"perspective\"][\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "sbert_encoder = SentenceTransformer(\"paraphrase-multilingual-mpnet-base-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus_embeddings = sbert_encoder.encode(corpus[\"argument\"].values)\n",
    "corpus[\"sbert_embeddings\"] = list(corpus_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_train_embeddings = sbert_encoder.encode(baseline_queries_train[\"text\"].values) \n",
    "query_dev_embeddings = sbert_encoder.encode(baseline_queries_dev[\"text\"].values) \n",
    "\n",
    "baseline_queries_train[\"sbert_embeddings\"] = list(query_train_embeddings)\n",
    "baseline_queries_dev[\"sbert_embeddings\"] = list(query_dev_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment when using test set\n",
    "\n",
    "# query_test_embeddings = sbert_encoder.encode(perspective_queries_test[\"text\"].values)\n",
    "\n",
    "# perspective_queries_test[\"sbert_embeddings\"] = list(query_test_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_similarities = cosine_similarity(\n",
    "    list(baseline_queries_train[\"sbert_embeddings\"].values), list(corpus[\"sbert_embeddings\"].values)\n",
    ")\n",
    "dev_similarities = cosine_similarity(\n",
    "    list(baseline_queries_dev[\"sbert_embeddings\"].values), list(corpus[\"sbert_embeddings\"].values)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment when using test set\n",
    "\n",
    "# test_similarities = cosine_similarity(\n",
    "#     list(perspective_queries_test[\"sbert_embeddings\"].values), list(corpus[\"sbert_embeddings\"].values)\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_scores = []\n",
    "for i in range(len(baseline_queries_train[\"query_id\"].values)):\n",
    "    row = train_similarities[i, :]\n",
    "    top_indices = np.argsort(row)[-50:]\n",
    "    topics = {}\n",
    "    for argument_index in top_indices:\n",
    "        current_topic = corpus.iloc[argument_index]['topic']\n",
    "        if current_topic in topics:\n",
    "            topics[current_topic] += 1\n",
    "        else:\n",
    "            topics[current_topic] = 1\n",
    "    train_scores.append({\n",
    "        'query_id': baseline_queries_train[\"query_id\"].values[i],\n",
    "        'similarity_scores': train_similarities[i],\n",
    "        'topic': topics\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_scores_df = pd.DataFrame(train_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_scores = []\n",
    "for i in range(len(baseline_queries_dev[\"query_id\"].values)):\n",
    "    row = dev_similarities[i, :]\n",
    "    top_indices = np.argsort(row)[-50:]\n",
    "    topics = {}\n",
    "    for argument_index in top_indices:\n",
    "        current_topic = corpus.iloc[argument_index]['topic']\n",
    "        if current_topic in topics:\n",
    "            topics[current_topic] += 1\n",
    "        else:\n",
    "            topics[current_topic] = 1\n",
    "    dev_scores.append({\n",
    "        'query_id': baseline_queries_dev[\"query_id\"].values[i],\n",
    "        'similarity_scores': dev_similarities[i],\n",
    "        'topic': topics\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_scores_df = pd.DataFrame(dev_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment when using test set\n",
    "# test_scores = []\n",
    "# for i in range(len(perspective_queries_test[\"query_id\"].values)):\n",
    "#     row = test_similarities[i, :]\n",
    "#     top_indices = np.argsort(row)[-50:]\n",
    "#     topics = {}\n",
    "#     for argument_index in top_indices:\n",
    "#         current_topic = corpus.iloc[argument_index]['topic']\n",
    "#         if current_topic in topics:\n",
    "#             topics[current_topic] += 1\n",
    "#         else:\n",
    "#             topics[current_topic] = 1\n",
    "#     test_scores.append({\n",
    "#         'query_id': perspective_queries_test[\"query_id\"].values[i],\n",
    "#         'similarity_scores': test_similarities[i],\n",
    "#         'topic': topics\n",
    "#     })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_scores_df = pd.DataFrame(test_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "destination_folder = \"final-scores\" # change destination folder name!\n",
    "\n",
    "train_scores_df.to_json(f\"{destination_folder}/baseline_train_similarity_scores.jsonl\", orient=\"records\", lines=True)\n",
    "dev_scores_df.to_json(f\"{destination_folder}/baseline_dev_similarity_scores.jsonl\", orient=\"records\", lines=True)\n",
    "\n",
    "# uncomment when using test set\n",
    "# test_scores_df.to_json(f\"{destination_folder}/surprise/perspective_similarity_scores.jsonl\", orient=\"records\", lines=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}