{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "import numpy as np\n",
    "import sklearn\n",
    "import sklearn.metrics as skm\n",
    "import sklearn.ensemble as ensemble\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import trange, tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = 'data/halomi_full.tsv'\n",
    "data_df = pd.read_csv(file, sep=\"\\t\")\n",
    "all_metrics = [\"score_log_loss\",\"score_alti_mean\",\"score_alti_t_mean\",\"score_attn_ot\",\"score_comet_qe\",\"score_labse\",\"score_laser\",\"score_xnli\"]\n",
    "# inerse all metrics\n",
    "for metric in all_metrics:\n",
    "    data_df[metric] = - data_df[metric]\n",
    "\n",
    "data_df[\"is_hall\"] = data_df[\"class_hall\"].apply(lambda x:1 if x[0]!='1' else 0)\n",
    "data_df[\"is_omit\"] = data_df[\"class_omit\"].apply(lambda x:1 if x[0]!='1' else 0)\n",
    "\n",
    "high_resource_languages = [\"arb_Arab\", \"deu_Latn\", \"rus_Cyrl\", \"spa_Latn\", \"zho_Hans\"]\n",
    "low_resource_language_pairs = [\"kas_Deva\", \"mni_Beng\", \"yor_Latn\"]\n",
    "\n",
    "data_high = data_df[\n",
    "    data_df[\"direction\"].apply(\n",
    "        lambda x: any([lang in x for lang in high_resource_languages])\n",
    "        and \"eng_Latn\" in x\n",
    "    )\n",
    "].copy()\n",
    "\n",
    "data_low = data_df[\n",
    "    data_df[\"direction\"].apply(\n",
    "        lambda x: any([lang in x for lang in low_resource_language_pairs])\n",
    "        and \"eng_Latn\" in x\n",
    "    )\n",
    "].copy()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Non hallucinations : 1919\n",
      "Hallucinations : 1010\n",
      "Non omisions : 1292\n",
      "Omisions : 1637\n"
     ]
    }
   ],
   "source": [
    "print(\"Non hallucinations :\",len(data_high[data_high[\"is_hall\"]==0]))\n",
    "print(\"Hallucinations :\",len(data_high[data_high[\"is_hall\"]==1]))\n",
    "\n",
    "print(\"Non omisions :\",len(data_high[data_high[\"is_omit\"]==0]))\n",
    "print(\"Omisions :\",len(data_high[data_high[\"is_omit\"]==1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hallucination Detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def auroc_fpr(pred, num_thresh = 100):\n",
    "    thrs = np.linspace(start=pred.min(), stop=pred.max(), num=num_thresh)\n",
    "    decision_by_thr = np.zeros((len(pred), num_thresh))\n",
    "    \n",
    "    for i in range(len(thrs)):\n",
    "        thr = thrs[i]\n",
    "        y_pred = np.where(pred <= thr, 1, 0)\n",
    "        decision_by_thr[:, i] = y_pred\n",
    "    return decision_by_thr\n",
    "\n",
    "def calculate_auc_fp_from_pred(pred, testdata_is_hall):\n",
    "    decision_by_thr = auroc_fpr(pred, num_thresh = 100)\n",
    "    is_hall = np.zeros(decision_by_thr.shape)\n",
    "    for i in range(is_hall.shape[1]):\n",
    "        is_hall[:, i] = testdata_is_hall.to_numpy()\n",
    "\n",
    "    tp = np.where((decision_by_thr == 1) & (is_hall==1), 1, 0)\n",
    "    fn = np.where((decision_by_thr == 0) & (is_hall==1), 1, 0)\n",
    "    tn = np.where((decision_by_thr == 0) & (is_hall==0), 1, 0)\n",
    "    fp = np.where((decision_by_thr == 1) & (is_hall==0), 1, 0)\n",
    "\n",
    "    tpr = tp.sum(axis=0) / (tp.sum(axis=0) + fn.sum(axis=0))\n",
    "    fpr = fp.sum(axis=0) / (fp.sum(axis=0) + tn.sum(axis=0))\n",
    "\n",
    "    x = np.interp(0.90, np.sort(tpr), np.sort(fpr))\n",
    "    auc, fp = skm.auc(np.sort(fpr), np.sort(tpr)) * 100, x * 100\n",
    "    return auc, fp\n",
    "\n",
    "def calculate_auc_fp(trainset_data_detectors,testset_data_detectors,testdata_is_hall,prediction_method,num_try = 10,*args, **kwargs):\n",
    "    all_auc, all_fp = 0, 0\n",
    "    for _ in range(num_try):\n",
    "        pred = prediction_method(trainset_data_detectors,testset_data_detectors,*args, **kwargs)\n",
    "        auc, fp = calculate_auc_fp_from_pred(pred, testdata_is_hall)\n",
    "        all_auc += auc\n",
    "        all_fp += fp\n",
    "    return all_auc / num_try, all_fp / num_try, pred\n",
    "\n",
    "def get_best_threshold(pred, testdata_is_hall,num_thresh = 100):\n",
    "    thrs = np.linspace(start=pred.min(), stop=pred.max(), num=num_thresh)\n",
    "    decision_by_thr = auroc_fpr(pred, num_thresh = num_thresh)\n",
    "    is_hall = np.zeros(decision_by_thr.shape)\n",
    "    for i in range(is_hall.shape[1]):\n",
    "        is_hall[:, i] = testdata_is_hall.to_numpy()\n",
    "    f1 = np.zeros(decision_by_thr.shape[1])\n",
    "    for i in range(is_hall.shape[1]):\n",
    "        f1[i] = skm.f1_score(is_hall[:, i], decision_by_thr[:, i])\n",
    "    return thrs[np.argmax(f1)], f1.max()\n",
    "\n",
    "def plot_prediction(pred, testdata_is_hall, title, thr, nbins=50,legend=True):\n",
    "    plt.figure(figsize=(10, 3))\n",
    "    bins = np.histogram(pred, bins=nbins)[1]\n",
    "\n",
    "    sns.histplot(pred[testdata_is_hall == 0], bins=bins, color='green', label=\"No Hallucination\", kde=True, edgecolor='white', stat='percent')\n",
    "    sns.histplot(pred[testdata_is_hall == 1], bins=bins, color='red', label=\"Hallucination\", kde=True, edgecolor='white', stat='percent')\n",
    "    # plt.axvline(x=thr, color='black', linestyle='--', linewidth=2, label='Threshold')\n",
    "\n",
    "    plt.xlabel('Score', fontsize=30)\n",
    "    plt.ylabel('Proportion', fontsize=30)\n",
    "    # plt.title(title, fontsize=16)\n",
    "    sns.set_palette(\"Set2\")\n",
    "    if legend:\n",
    "        plt.legend(loc='upper left', fontsize=30)\n",
    "    plt.grid(True, linestyle='--', alpha=0.7)\n",
    "    sns.despine()\n",
    "    plt.xticks(fontsize=25)\n",
    "    plt.yticks(fontsize=25)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prediction_method_single(trainset_data_detectors,testset_data_detectors): # single detector\n",
    "    pred = testset_data_detectors\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prediction_method_IF(trainset_data_detectors,testset_data_detectors): # aggregation of detectors using isolation forest\n",
    "    isolation_forest = ensemble.IsolationForest(n_estimators=500,n_jobs=-1)\n",
    "    clf = isolation_forest.fit(trainset_data_detectors)\n",
    "    pred = clf.score_samples(testset_data_detectors)\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "def prediction_logistic_regression(trainset_data_detectors,testset_data_detectors,trainset_data_is_hall):\n",
    "    logistic_regression = LogisticRegression(max_iter=1000,solver='liblinear')\n",
    "    trainset_data_is_hall = np.array(trainset_data_is_hall).ravel()\n",
    "    logistic_regression.fit(trainset_data_detectors, trainset_data_is_hall)\n",
    "    pred = logistic_regression.predict_proba(testset_data_detectors)[:, 1]\n",
    "    return -pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = {\"high\":data_high,\"low\":data_low,\"full\":data_df}\n",
    "oods = [\"is_hall\", \"is_omit\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_dfs = 10\n",
    "p = 0.1 # fraction of data to use for training\n",
    "data_df_trainsets = {}\n",
    "data_df_testsets = {}\n",
    "for dataset in datasets:\n",
    "    data_df_trainsets[dataset] = []\n",
    "    data_df_testsets[dataset] = []\n",
    "    for i in range(N_dfs):\n",
    "        data_df_trainsets[dataset].append(datasets[dataset].sample(frac=p, random_state=i))\n",
    "        data_df_testsets[dataset].append(datasets[dataset].drop(data_df_trainsets[dataset][i].index))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_metrics = [\"score_log_loss\", \"score_alti_mean\", \"score_attn_ot\"]\n",
    "external_metrics = [\"score_comet_qe\",\"score_labse\",\"score_laser\",\"score_xnli\"]\n",
    "model_external_metrics = model_metrics + external_metrics\n",
    "sorted_metrics = [\"score_labse\",\"score_log_loss\",\"score_xnli\",\"score_comet_qe\",\"score_alti_mean\",\"score_laser\",\"score_attn_ot\"]\n",
    "\n",
    "metrics_type = {\n",
    "    \"model\":model_metrics,\n",
    "    \"external\":external_metrics,\n",
    "    \"all\":model_external_metrics}\n",
    "\n",
    "\n",
    "normalized_metrics = {}\n",
    "normalized_metrics_trainset = {}\n",
    "for dataset in datasets:\n",
    "    normalized_metrics[dataset] = {}\n",
    "    normalized_metrics_trainset[dataset] = {}\n",
    "    for i in range(N_dfs):\n",
    "        normalized_metrics[dataset][i] = {}\n",
    "        normalized_metrics_trainset[dataset][i] = {}\n",
    "        for metric in model_external_metrics:\n",
    "            normalized_metrics[dataset][i][metric] = data_df_testsets[dataset][i][metric].apply(lambda x: (x - data_df_trainsets[dataset][i][metric].min()) / (data_df_trainsets[dataset][i][metric].max() - data_df_trainsets[dataset][i][metric].min()))\n",
    "            normalized_metrics_trainset[dataset][i][metric] = data_df_trainsets[dataset][i][metric].apply(lambda x: (x - data_df_trainsets[dataset][i][metric].min()) / (data_df_trainsets[dataset][i][metric].max() - data_df_trainsets[dataset][i][metric].min()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configs = {}\n",
    "# for dataset in datasets:\n",
    "for dataset in [\"high\"]:\n",
    "    configs[dataset] = {}\n",
    "    for i in range(N_dfs):\n",
    "        data_df_testset = data_df_testsets[dataset][i]\n",
    "        data_df_trainset = data_df_trainsets[dataset][i]\n",
    "        configs[dataset][i] = []\n",
    "        # for ood in oods:\n",
    "        for ood in [\"is_hall\"]:\n",
    "            for metric in model_external_metrics:\n",
    "                config = {\"ood\":ood,\"name\":metric,\"trainset_data_detectors\":None,\"testset_data_detectors\":data_df_testset[metric],\n",
    "                        \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "                        \"trainset_data_is_hall\":None, \"method\":None}\n",
    "                configs[dataset][i].append(config)\n",
    "            # Sum_norm\n",
    "            for name,metrics in metrics_type.items():\n",
    "                config = {\"ood\":ood,\"name\":\"Sum_norm_\"+name,\"trainset_data_detectors\":None,\n",
    "                        \"testset_data_detectors\":np.sum(pd.DataFrame(normalized_metrics[dataset][i])[metrics],axis=1),\n",
    "                        \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "                        \"trainset_data_is_hall\":None, \"method\":None}\n",
    "                configs[dataset][i].append(config)\n",
    "            # Max_norm\n",
    "            for name,metrics in metrics_type.items():\n",
    "                config = {\"ood\":ood,\"name\":\"Max_norm_\"+name,\"trainset_data_detectors\":None,\n",
    "                        \"testset_data_detectors\":np.max(pd.DataFrame(normalized_metrics[dataset][i])[metrics],axis=1),\n",
    "                        \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "                        \"trainset_data_is_hall\":None, \"method\":None}\n",
    "                configs[dataset][i].append(config)\n",
    "            # IF\n",
    "            for name,metrics in metrics_type.items():\n",
    "                config = {\"ood\":ood,\"name\":\"IF_\"+name,\"trainset_data_detectors\":data_df_trainset[metrics],\n",
    "                        \"testset_data_detectors\":data_df_testset[metrics],\n",
    "                        \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_IF,\"num_try\":1,\n",
    "                        \"trainset_data_is_hall\":None, \"method\":None}\n",
    "                configs[dataset][i].append(config)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_full = {}\n",
    "for dataset in datasets:\n",
    "    results_full[dataset] = {}\n",
    "    for config in configs[dataset][0]:\n",
    "        if config[\"ood\"] not in results_full[dataset]:\n",
    "            results_full[dataset][config[\"ood\"]] = {}\n",
    "        results_full[dataset][config[\"ood\"]][config[\"name\"]] = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dataset in tqdm(configs):\n",
    "    for i in trange(N_dfs):\n",
    "        for config in configs[dataset][i]:\n",
    "            if config[\"ood\"] not in results_full[dataset] or len(results_full[dataset][config[\"ood\"]][config[\"name\"]]) < i+1:\n",
    "                if config[\"trainset_data_is_hall\"] is None:\n",
    "                    if config[\"method\"] is None:\n",
    "                        auc, fp, pred = calculate_auc_fp(config[\"trainset_data_detectors\"],config[\"testset_data_detectors\"], config[\"testdata_is_hall\"], \n",
    "                                                    config[\"prediction_method\"], num_try = config[\"num_try\"])\n",
    "                    else:\n",
    "                        auc, fp, pred = calculate_auc_fp(config[\"trainset_data_detectors\"],config[\"testset_data_detectors\"], config[\"testdata_is_hall\"], \n",
    "                                                    config[\"prediction_method\"], num_try = config[\"num_try\"], method = config[\"method\"])\n",
    "                    \n",
    "                else:\n",
    "                    auc, fp, pred = calculate_auc_fp(config[\"trainset_data_detectors\"],config[\"testset_data_detectors\"], config[\"testdata_is_hall\"], \n",
    "                                                config[\"prediction_method\"], num_try = config[\"num_try\"],\n",
    "                                                trainset_data_is_hall =config[\"trainset_data_is_hall\"])\n",
    "                thr, f1 = get_best_threshold(pred, config[\"testdata_is_hall\"])\n",
    "                results_full[dataset][config[\"ood\"]][config[\"name\"]].append({\"auc\":auc,\"fp\":fp, \"f1_score\":f1, \"pred\":pred, \"thr\":thr})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save results\n",
    "# import pickle\n",
    "# with open('results_full_halomi.pickle', 'wb') as handle:\n",
    "#     pickle.dump(results_full, handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load results\n",
    "import pickle\n",
    "with open('results_full_halomi.pickle', 'rb') as handle:\n",
    "    results_full = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 81600 characters]",
      "text/plain": [
       "<Figure size 1000x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 75112 characters]",
      "text/plain": [
       "<Figure size 1000x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "m_to_plot = \"score_labse\"\n",
    "ood = \"is_hall\"\n",
    "dataset = \"high\"\n",
    "i=0\n",
    "plot_prediction(results_full[dataset][ood][m_to_plot][i][\"pred\"],data_df_testsets[dataset][i]['is_hall'].values,\"\",results_full[dataset][ood][m_to_plot][i][\"thr\"],nbins=100)\n",
    "m_to_plot = \"Sum_norm_all\"\n",
    "plot_prediction(results_full[dataset][ood][m_to_plot][i][\"pred\"]/len(model_external_metrics),data_df_testsets[dataset][i]['is_hall'].values,\"\",results_full[dataset][ood][m_to_plot][i][\"thr\"]/len(model_external_metrics),nbins=100)\n"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
