{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5d31096e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import math\n",
    "import openai\n",
    "import textstat\n",
    "import tiktoken\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from rouge_score import rouge_scorer\n",
    "from nltk.stem import WordNetLemmatizer\n",
    "from scipy.spatial.distance import cosine\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c237d94b",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = 'Hemonc'\n",
    "dataset = pd.read_csv(f'Data/Input/{data}.csv')\n",
    "features = pd.DataFrame()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9b2d8cb",
   "metadata": {},
   "source": [
    "### Complexity\n",
    "\n",
    "- Length\n",
    "- Readability\n",
    "- Number of Unique Tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b0c1e3ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "wnl = WordNetLemmatizer()\n",
    "get_length = lambda x: len(x)\n",
    "get_readability = lambda x: textstat.flesch_kincaid_grade(x)\n",
    "count_unitok = lambda x: len(set([wnl.lemmatize(tkn).lower() for tkn in x.split(' ')]))\n",
    "\n",
    "features['Length'] = dataset[f'evidence {augmentation}'].apply(get_length)\n",
    "features['Readability'] = dataset[f'evidence {augmentation}'].apply(get_readability)\n",
    "features['# Uniq Tkns'] = dataset[f'evidence {augmentation}'].apply(count_unitok)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ff9d84b",
   "metadata": {},
   "source": [
    "### Relevance\n",
    "\n",
    "- N-gram Overlap (ROUGE)\n",
    "- Embedding Similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "79b2b273",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "chunking: 100%|█████████████████████████████████████████████████| 6212/6212 [00:04<00:00, 1513.14it/s]\n",
      "embedding: 100%|█████████████████████████████████████████████████| 6212/6212 [00:25<00:00, 246.59it/s]\n",
      "embedding: 100%|█████████████████████████████████████████████████| 6212/6212 [00:20<00:00, 304.69it/s]\n"
     ]
    }
   ],
   "source": [
    "emb_model, chunk_size = \"text-embedding-3-large\", 8000\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(model_name=emb_model, chunk_size=chunk_size)\n",
    "evidences, questions, idx2pieces = [], [], {}\n",
    "for idx, row in tqdm(dataset.iterrows(), total=len(dataset), desc='chunking'):\n",
    "    evidence = text_splitter.split_text(row[f'evidence {augmentation}'])\n",
    "    idx2pieces[idx] = (len(idx2pieces), len(idx2pieces)+len(evidence))\n",
    "    evidences += evidence; questions.append(row['question 1'])\n",
    "\n",
    "encoding = tiktoken.encoding_for_model(emb_model)\n",
    "client = openai.OpenAI(api_key=os.environ[\"OPENAI_API_KEY\"])\n",
    "batch_limit, query_limit = 500000, 1000\n",
    "evi_embs_, qn_embs = [], []\n",
    "for texts, embs in zip([evidences, questions], [evi_embs_, qn_embs]):\n",
    "    batch, batch_size = [], 0\n",
    "    for text in tqdm(texts, desc='embedding'):\n",
    "        text_size = len(encoding.encode(text))\n",
    "        if batch_size + text_size > batch_limit or len(batch) >= query_limit:\n",
    "            response = client.embeddings.create(input=batch, model=emb_model)\n",
    "            embs += [each.embedding for each in response.data]\n",
    "            batch, batch_size = [], 0\n",
    "        batch.append(text); batch_size += text_size\n",
    "    response = client.embeddings.create(input=batch, model=emb_model)\n",
    "    embs += [np.array(each.embedding) for each in response.data]\n",
    "    \n",
    "evi_embs = []\n",
    "for idx, (start, end) in idx2pieces.items():\n",
    "    evi_embs.append(np.mean(evi_embs_[start:end], axis=0))\n",
    "for idx, (evi_emb, qn_emb) in enumerate(zip(evi_embs, qn_embs)):\n",
    "    features.at[idx, 'Emb Sim'] = 1 - cosine(evi_emb, qn_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d2840f61",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████| 6212/6212 [00:09<00:00, 624.96it/s]\n"
     ]
    }
   ],
   "source": [
    "scorer = rouge_scorer.RougeScorer(['rouge2'], use_stemmer=True)\n",
    "for m in ['P', 'R', 'F']:\n",
    "    features[f'rouge2-{m}'] = None\n",
    "\n",
    "for idx, row in tqdm(dataset.iterrows(), total=len(dataset)):\n",
    "    scores = scorer.score(row['question 1'], row[f'evidence {augmentation}'])\n",
    "    features.at[idx, f'rouge2-P'] = scores[f'rouge2'].precision\n",
    "    features.at[idx, f'rouge2-R'] = scores[f'rouge2'].recall\n",
    "    features.at[idx, f'rouge2-F'] = scores[f'rouge2'].fmeasure"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecee46ad",
   "metadata": {},
   "source": [
    "### Familiarity\n",
    "\n",
    "- Perplexity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "6352a53b",
   "metadata": {},
   "outputs": [],
   "source": [
    "abbre2model = {'Gemma-2B':'gemma-2-2b-it', 'Gemma-9B':'gemma-2-9b-it', 'Gemma-27B':'gemma-2-27b-it', \n",
    "               'Llama-3B':'Llama-3.2-3B-Instruct', 'Llama-8B':'Llama-3.1-8B-Instruct', 'Llama-70B':'Llama-3.3-70B-Instruct', \n",
    "               'Qwen-3B':'Qwen2.5-3B-Instruct', 'Qwen-7B':'Qwen2.5-7B-Instruct', 'Qwen-14B':'Qwen2.5-14B-Instruct'}\n",
    "for abbre, model in abbre2model.items():\n",
    "    for fea in ['Perplexity', 'Entropy']:\n",
    "        raw = json.load(open(f'Data/Uncertainty/{model}_{fea}.json'))\n",
    "        features[f'{abbre} {fea[:4]} Context'] = raw[data]['evidence']\n",
    "        features[f'{abbre} {fea[:4]} Question'] = raw[data]['question 1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "7bb4cfd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "features.to_csv(f'Data/Feature/{data}.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
