05_cook_artifacts.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div align='left' style=\"width:29%;overflow:hidden;\">\n",
    "<a href='http://inria.fr'>\n",
    "<img src='https://github.com/lmarti/jupyter_custom/raw/master/imgs/inr_logo_rouge.png' alt='Inria logo' title='Inria logo'/>\n",
    "</a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cooking RISOTTO artifacts\n",
    "\n",
    "> This notebook purpose is to build the binary files used by RISOTTO's GUI.\n",
    "\n",
    "We'll build three Pandas DataFrames:\n",
    "\n",
    "- `papers`:\n",
    "contains the data of the papers, including the `cord_uid` identifier and the PageRank scores.\n",
    "\n",
    "- `papers_topics`:\n",
    "contains the association between papers, topics, and subtopics.\n",
    "The papers are indexed by their `cord_uid` identifier.\n",
    "\n",
    "- `topics`:\n",
    "contains the token pseudocounts of the different topics and subtopics.\n",
    "\n",
    "Each one of the previously defined DataFrames will be stored in a single HDF file named `artifacts.hdf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all_flag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from pathlib import Path\n",
    "\n",
    "from fastprogress.fastprogress import progress_bar\n",
    "import pandas as pd\n",
    "import networkx as nx\n",
    "\n",
    "from risotto.references import load_papers_from_metadata_file, process_references, build_papers_reference_graph\n",
    "from risotto.lda import process_papers_file_contents\n",
    "from risotto.hierarchical_lda import fit_lda_model\n",
    "from risotto.zero_shot import embed_papers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CORD19_DATASET_FOLDER = Path(\"./datasets/CORD-19-research-challenge\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "ARTIFACTS_FOLDER_NAME = \"artifacts\"\n",
    "ARTIFACTS_DEFAULT_PATH = f\"{ARTIFACTS_FOLDER_NAME}/artifacts.hdf\"\n",
    "\n",
    "\n",
    "def remove_duplicated_idxs(df):\n",
    "    df = df.loc[~df.index.duplicated(keep='first')]\n",
    "    return df\n",
    "\n",
    "\n",
    "def load_papers_metadata(dataset_folder):\n",
    "    metadata_df = pd.read_csv(f\"{dataset_folder}/metadata.csv\",\n",
    "                              index_col=\"cord_uid\")\n",
    "    return metadata_df\n",
    "\n",
    "\n",
    "def _get_affiliation_country(author):\n",
    "    affiliation = author[\"affiliation\"]\n",
    "    if len(affiliation) == 0:\n",
    "        return None, None\n",
    "    # Country\n",
    "    country = affiliation[\"location\"][\"country\"] if \"country\" in affiliation[\n",
    "        \"location\"] else None\n",
    "    # Affiliation\n",
    "    laboratory = affiliation[\"laboratory\"]\n",
    "    institution = affiliation[\"institution\"] if len(\n",
    "        affiliation[\"institution\"]) > 0 else None\n",
    "    if institution is not None and len(laboratory) > 0:\n",
    "        institution += f\" ({laboratory})\"\n",
    "    return institution, country\n",
    "\n",
    "\n",
    "def build_papers_artifact(dataset_folder,\n",
    "                          should_dump=True,\n",
    "                          dump_path=ARTIFACTS_DEFAULT_PATH):\n",
    "    \"\"\"\n",
    "    Returns:\n",
    "        - `joined_df`: a DataFrame with the dataset metadata columns and the PageRank score\n",
    "            of each paper. The rows are indexed by the `cord_uid` identifier.\n",
    "    \"\"\"\n",
    "    metadata_df = remove_duplicated_idxs(load_papers_metadata(dataset_folder))\n",
    "    papers = load_papers_from_metadata_file(dataset_folder)\n",
    "    process_references(papers)\n",
    "    references_graph = build_papers_reference_graph(papers)\n",
    "    pagerank = nx.pagerank(references_graph)\n",
    "\n",
    "    # Build PageRank DataFrame dict\n",
    "    df_dict = {\n",
    "        \"cord_uid\": [],\n",
    "        \"pagerank\": [],\n",
    "        \"affiliation\": [],\n",
    "        \"country\": []\n",
    "    }\n",
    "    for paper, score in pagerank.items():\n",
    "        # PageRank\n",
    "        df_dict[\"cord_uid\"].append(paper._metadata_row.name)\n",
    "        df_dict[\"pagerank\"].append(score)\n",
    "\n",
    "        # Affiliations and countries\n",
    "        affiliations = set()\n",
    "        countries = set()\n",
    "        for author in paper._file_contents[\"metadata\"][\"authors\"]:\n",
    "            affiliation, country = _get_affiliation_country(author)\n",
    "            if affiliation is not None:\n",
    "                affiliations.add(affiliation)\n",
    "            if country is not None:\n",
    "                countries.add(country)\n",
    "        df_dict[\"affiliation\"].append(\", \".join(affiliations))\n",
    "        df_dict[\"country\"].append(\", \".join(countries))\n",
    "\n",
    "    pagerank_df = pd.DataFrame.from_dict(df_dict)\n",
    "    pagerank_df = remove_duplicated_idxs(pagerank_df.set_index(\"cord_uid\"))\n",
    "\n",
    "    joined_df = pagerank_df.join(metadata_df)\n",
    "\n",
    "    if should_dump:\n",
    "        joined_df.to_hdf(dump_path, key=\"papers\")\n",
    "\n",
    "    return joined_df, pagerank_df, metadata_df\n",
    "\n",
    "\n",
    "def load_papers_artifact(artifacts_path=ARTIFACTS_DEFAULT_PATH):\n",
    "    return pd.read_hdf(artifacts_path, key=\"papers\")\n",
    "\n",
    "\n",
    "def build_papers_topics_artifact(dataset_folder,\n",
    "                                 should_dump=True,\n",
    "                                 dump_path=ARTIFACTS_DEFAULT_PATH):\n",
    "    \"\"\"\n",
    "    Returns:\n",
    "        - `papers_df`: a DataFrame. Each row represents a paper and has the following columns:\n",
    "            - `cord_uid`: index\n",
    "            - `text`: paper text\n",
    "            - `topic`: first-level topic identifier\n",
    "            - `subtopic`: second-level topic identifier\n",
    "        - `topics_df`: a DataFrame. Each row represents a token and has the following columns:\n",
    "            - `token`: index\n",
    "            - `{topic_id}`: token pseudocount in first-level topic {topic_id}\n",
    "            - `{topic_id}-{subtopic_id}`: token pseudocount in second-level {subtopic_id} of first-level {topic_id}.\n",
    "        NaNs values mean that the token wasn't considered in the topic modelling step.\n",
    "    \"\"\"\n",
    "    # Load the papers contents\n",
    "    papers = load_papers_from_metadata_file(dataset_folder)\n",
    "    docs = process_papers_file_contents(papers)\n",
    "\n",
    "    # Here we build a dataframe with the text of each dataset\n",
    "    # indexed by its own `cord_uid`\n",
    "    papers_df_dict = {\"cord_uid\": []}\n",
    "    for paper in papers:\n",
    "        papers_df_dict[\"cord_uid\"].append(paper._metadata_row.name)\n",
    "    papers_df_dict[\"text\"] = docs\n",
    "    papers_df = pd.DataFrame.from_dict(papers_df_dict)\n",
    "    papers_df = papers_df.set_index(\"cord_uid\")\n",
    "\n",
    "    # Fit the first level LDA model and classify the papers\n",
    "    lda, docs_vectorized, vectorizer = fit_lda_model(\n",
    "        papers_df[\"text\"].tolist(),\n",
    "        n_components=8,\n",
    "        n_jobs=4,\n",
    "    )\n",
    "    docs_classified = lda.transform(docs_vectorized)\n",
    "    docs_topics = docs_classified.argmax(1)\n",
    "    papers_df[\"topic\"] = docs_topics\n",
    "\n",
    "    # Fit the second level LDA models and classify the papers\n",
    "    sub_ldas = {}\n",
    "    sub_docs_vectorized = {}\n",
    "    sub_docs_topics = {}\n",
    "    sub_vectorizers = {}\n",
    "    for topic_id, group_df in progress_bar(papers_df.groupby(by=\"topic\")):\n",
    "        _sub_lda, _sub_docs_vectorized, _sub_vectorizer = fit_lda_model(\n",
    "            group_df[\"text\"].tolist(),\n",
    "            n_components=4,\n",
    "            n_jobs=4,\n",
    "        )\n",
    "        sub_ldas[topic_id] = _sub_lda\n",
    "        sub_docs_vectorized[topic_id] = _sub_docs_vectorized\n",
    "        group_topics = _sub_lda.transform(_sub_docs_vectorized).argmax(1)\n",
    "        sub_docs_topics[topic_id] = pd.Series(group_topics,\n",
    "                                              index=group_df.index)\n",
    "        sub_vectorizers[topic_id] = _sub_vectorizer\n",
    "\n",
    "    papers_subtopics_list = list(sub_docs_topics.values())\n",
    "    papers_subtopics = remove_duplicated_idxs(\n",
    "        pd.concat(papers_subtopics_list).rename(\"subtopic\"))\n",
    "    papers_df = remove_duplicated_idxs(papers_df.join(papers_subtopics))\n",
    "\n",
    "    # Next step is to build a DataFrame with each topic pseudocount\n",
    "    # of its tokens.\n",
    "    topics_df = pd.DataFrame(\n",
    "        index=pd.Series(vectorizer.get_feature_names(), name=\"token\"))\n",
    "    for topic_idx, topic in enumerate(lda.components_):\n",
    "        topics_df[f\"{topic_idx}\"] = topic\n",
    "\n",
    "    subtopics_dfs = {}\n",
    "    for topic_idx, sub_lda in sub_ldas.items():\n",
    "        subtopic_df = pd.DataFrame(index=pd.Series(\n",
    "            sub_vectorizers[topic_idx].get_feature_names(), name=\"token\"))\n",
    "        for subtopic_idx, subtopic in enumerate(sub_lda.components_):\n",
    "            subtopic_df[f\"{topic_idx}-{subtopic_idx}\"] = subtopic\n",
    "        subtopics_dfs[topic_idx] = subtopic_df\n",
    "\n",
    "    for subtopic_df in subtopics_dfs.values():\n",
    "        topics_df = topics_df.join(subtopic_df, how=\"outer\")\n",
    "\n",
    "    if should_dump:\n",
    "        papers_df.to_hdf(dump_path, key=\"papers_topics\")\n",
    "        topics_df.to_hdf(dump_path, key=\"topics\")\n",
    "\n",
    "    return papers_df, topics_df\n",
    "\n",
    "\n",
    "def load_papers_topics_artifacts(artifacts_path=ARTIFACTS_DEFAULT_PATH):\n",
    "    return pd.read_hdf(artifacts_path, key=\"papers_topics\")\n",
    "\n",
    "\n",
    "def load_topics_artifacts(artifacts_path=ARTIFACTS_DEFAULT_PATH):\n",
    "    return pd.read_hdf(artifacts_path, key=\"topics\")\n",
    "\n",
    "\n",
    "def build_artifacts(dataset_folder, should_dump=True):\n",
    "    metapagerank_df, _, _ = build_papers_artifact(dataset_folder, should_dump)\n",
    "    papers_df, topics_df = build_papers_topics_artifact(\n",
    "        dataset_folder, should_dump)\n",
    "    return metapagerank_df, papers_df, topics_df\n",
    "\n",
    "\n",
    "def build_papers_embeddings(should_dump=True,\n",
    "                            dump_path=ARTIFACTS_DEFAULT_PATH):\n",
    "    papers = load_papers_artifact()\n",
    "    embeddings = embed_papers(papers)\n",
    "\n",
    "    if should_dump:\n",
    "        embeddings.to_hdf(dump_path, key=\"papers_embeddings\")\n",
    "\n",
    "    return embeddings\n",
    "\n",
    "\n",
    "def load_papers_embeddings(artifacts_path=ARTIFACTS_DEFAULT_PATH):\n",
    "    return pd.read_hdf(artifacts_path, key=\"papers_embeddings\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metapagerank_df, papers_df, topics_df = build_artifacts(CORD19_DATASET_FOLDER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joined_df, pagerank_df, metadata_df = build_papers_artifact(CORD19_DATASET_FOLDER)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}