5_SparkNLP_pubmed_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division, print_function\n",
    "\n",
    "from Bio import Entrez, Medline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyspark\n",
    "from pyspark.ml import Pipeline, feature as spark_ft, classification as spark_cls\n",
    "from sklearn import metrics as skmetrics\n",
    "import wordcloud\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "packages = [\n",
    "    'com.databricks:spark-xml_2.11:0.4.1',\n",
    "    'JohnSnowLabs:spark-nlp:1.4.1'\n",
    "]\n",
    "\n",
    "spark = pyspark.sql.SparkSession.builder \\\n",
    "    .master('local[4]') \\\n",
    "    .appName('notebook') \\\n",
    "    .config('spark.jars', 'pysparknlp-1.0.0/lib/sparknlp.jar') \\\n",
    "    .config('spark.jars.packages', ','.join(packages)) \\\n",
    "    .getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sparknlp.annotator import *\n",
    "from sparknlp.common import *\n",
    "from sparknlp.base import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def query(terms, num_docs=1000):\n",
    "    search_term = '+'.join(terms)\n",
    "    print('Searching PubMed abstracts for documents containing term: ',search_term)\n",
    "    handle = Entrez.esearch(db=\"pubmed\", term=search_term, retmax=num_docs)\n",
    "    record = Entrez.read(handle)\n",
    "    handle.close()\n",
    "    idlist = record[\"IdList\"]\n",
    "    \n",
    "    handle = Entrez.efetch(db=\"pubmed\", id=idlist, rettype=\"medline\",retmode=\"text\")\n",
    "    records = Medline.parse(handle)\n",
    "    data = []\n",
    "    for record in records:\n",
    "        data.append((record.get(\"TI\", \"?\"),record.get(\"AU\", \"?\"),record.get(\"SO\", \"?\"),record.get(\"AB\",\"?\")))\n",
    "\n",
    "    df = pd.DataFrame(data=data, columns=['Title','Authors','Source','Abstract'])\n",
    "    df.head(10)\n",
    "\n",
    "    df.replace(r'^\\?$', np.nan, regex=True, inplace=True)\n",
    "    df['Authors'] = df['Authors'].apply(lambda x: x if isinstance(x, list) else [])\n",
    "    df.fillna('', inplace=True)\n",
    "    df['Topic'] = search_term\n",
    "    \n",
    "    return spark.createDataFrame(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "topics = [\n",
    "    ['type', '1', 'diabetes'], \n",
    "    ['creutzfeldt', 'jakob', 'disease'], \n",
    "    ['post', 'traumatic', 'stress', 'disorder'],\n",
    "    ['heart', 'disease']\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "terms ['type', '1', 'diabetes'] num_docs 710\n",
      "Searching PubMed abstracts for documents containing term:  type+1+diabetes\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.6/site-packages/Bio/Entrez/__init__.py:564: UserWarning: \n",
      "Email address is not specified.\n",
      "\n",
      "To make use of NCBI's E-utilities, NCBI requires you to specify your\n",
      "email address with each request.  As an example, if your email address\n",
      "is A.N.Other@example.com, you can specify it as follows:\n",
      "   from Bio import Entrez\n",
      "   Entrez.email = 'A.N.Other@example.com'\n",
      "In case of excessive usage of the E-utilities, NCBI will attempt to contact\n",
      "a user at the email address provided before blocking access to the\n",
      "E-utilities.\n",
      "  E-utilities.\"\"\", UserWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "terms ['creutzfeldt', 'jakob', 'disease'] num_docs 565\n",
      "Searching PubMed abstracts for documents containing term:  creutzfeldt+jakob+disease\n",
      "terms ['post', 'traumatic', 'stress', 'disorder'] num_docs 582\n",
      "Searching PubMed abstracts for documents containing term:  post+traumatic+stress+disorder\n",
      "terms ['heart', 'disease'] num_docs 522\n",
      "Searching PubMed abstracts for documents containing term:  heart+disease\n"
     ]
    }
   ],
   "source": [
    "texts = None\n",
    "\n",
    "np.random.seed(123)\n",
    "for terms in topics:\n",
    "    num_docs = np.random.randint(200, 1000)\n",
    "    print('terms', terms, 'num_docs', num_docs)\n",
    "    if texts is None:\n",
    "        texts = query(terms, num_docs)\n",
    "    else:\n",
    "        texts = texts.union(query(terms, num_docs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2379"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "texts.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+--------------------+--------------------+--------------------+---------------+\n",
      "|               Title|             Authors|              Source|            Abstract|          Topic|\n",
      "+--------------------+--------------------+--------------------+--------------------+---------------+\n",
      "|A luciferase immu...|[Ling Y, Jiang P,...|Clin Biochem. 201...|AIM: Luciferase i...|type+1+diabetes|\n",
      "|Type 1 diabetes m...|[Singh RM, Howart...|Mol Cell Biochem....|There is much evi...|type+1+diabetes|\n",
      "|Erratum. Validati...|[Sosenko JM, Skyl...|Diabetes Care. 20...|                    |type+1+diabetes|\n",
      "|Macrovascular dis...|[Bjornstad P, Don...|Lancet Diabetes E...|Cardiovascular di...|type+1+diabetes|\n",
      "|Insufficient evid...|[Brignardello-Pet...|J Am Dent Assoc. ...|                    |type+1+diabetes|\n",
      "|Genetic risk scor...|[Thomas NJ, Jones...|Lancet Diabetes E...|                    |type+1+diabetes|\n",
      "|Genetic risk scor...|[Leslie RD, Lernm...|Lancet Diabetes E...|                    |type+1+diabetes|\n",
      "|Association betwe...|[Ahola AJ, Forsbl...|Diabetes Res Clin...|AIMS: Depressive ...|type+1+diabetes|\n",
      "|Alpha-1 antitryps...|[Weir GC, Ehlers ...|Pediatr Diabetes....|OBJECTIVE: To det...|type+1+diabetes|\n",
      "|Considering Cultu...|[Rose M, Aronow L...|Curr Diab Rep. 20...|PURPOSE OF REVIEW...|type+1+diabetes|\n",
      "|Improved Murine-M...|[Racine JJ, Stewa...|Diabetes. 2018 Fe...|Improved mouse mo...|type+1+diabetes|\n",
      "|ANNALS EXPRESS: C...|[Shimizu I, Hiram...|Ann Clin Biochem....|BackgroundTo clar...|type+1+diabetes|\n",
      "|Prevalence of ear...|[Adar A, Shalitin...|Diabetes Metab Re...|BACKGROUND: The i...|type+1+diabetes|\n",
      "|The Effect of Aro...|[Jeon YD, Kang SH...|J Med Food. 2018 ...|The number of dia...|type+1+diabetes|\n",
      "|Human Subcutaneou...|[Rigla M, Pons B,...|Diabetes Technol ...|BACKGROUND: Subcu...|type+1+diabetes|\n",
      "|Evaluation of Pan...|[Naganawa M, Lim ...|Mol Imaging Biol....|PURPOSE: Previous...|type+1+diabetes|\n",
      "|BMX-001, a novel ...|[Bruni A, Pepper ...|Am J Transplant. ...|Islet transplanta...|type+1+diabetes|\n",
      "|Continuous Glucos...|[Wood A, O'Neal D...|Intern Med J. 201...|The advent of dev...|type+1+diabetes|\n",
      "|Factitious hypogl...|[Bauman V, Sturke...|Pediatr Diabetes....|BACKGROUND: Facti...|type+1+diabetes|\n",
      "|Intensive remote ...|[Gandrud L, Altan...|Pediatr Diabetes....|OBJECTIVE: While ...|type+1+diabetes|\n",
      "+--------------------+--------------------+--------------------+--------------------+---------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "texts.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "non_empty_texts = texts.where('Abstract != \"\"') \\\n",
    "    .withColumn('id', pyspark.sql.functions.monotonically_increasing_id())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2180"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "non_empty_texts.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+--------------------+--------------------+--------------------+---------------+---+\n",
      "|               Title|             Authors|              Source|            Abstract|          Topic| id|\n",
      "+--------------------+--------------------+--------------------+--------------------+---------------+---+\n",
      "|A luciferase immu...|[Ling Y, Jiang P,...|Clin Biochem. 201...|AIM: Luciferase i...|type+1+diabetes|  0|\n",
      "|Type 1 diabetes m...|[Singh RM, Howart...|Mol Cell Biochem....|There is much evi...|type+1+diabetes|  1|\n",
      "|Macrovascular dis...|[Bjornstad P, Don...|Lancet Diabetes E...|Cardiovascular di...|type+1+diabetes|  2|\n",
      "|Association betwe...|[Ahola AJ, Forsbl...|Diabetes Res Clin...|AIMS: Depressive ...|type+1+diabetes|  3|\n",
      "|Alpha-1 antitryps...|[Weir GC, Ehlers ...|Pediatr Diabetes....|OBJECTIVE: To det...|type+1+diabetes|  4|\n",
      "|Considering Cultu...|[Rose M, Aronow L...|Curr Diab Rep. 20...|PURPOSE OF REVIEW...|type+1+diabetes|  5|\n",
      "|Improved Murine-M...|[Racine JJ, Stewa...|Diabetes. 2018 Fe...|Improved mouse mo...|type+1+diabetes|  6|\n",
      "|ANNALS EXPRESS: C...|[Shimizu I, Hiram...|Ann Clin Biochem....|BackgroundTo clar...|type+1+diabetes|  7|\n",
      "|Prevalence of ear...|[Adar A, Shalitin...|Diabetes Metab Re...|BACKGROUND: The i...|type+1+diabetes|  8|\n",
      "|The Effect of Aro...|[Jeon YD, Kang SH...|J Med Food. 2018 ...|The number of dia...|type+1+diabetes|  9|\n",
      "|Human Subcutaneou...|[Rigla M, Pons B,...|Diabetes Technol ...|BACKGROUND: Subcu...|type+1+diabetes| 10|\n",
      "|Evaluation of Pan...|[Naganawa M, Lim ...|Mol Imaging Biol....|PURPOSE: Previous...|type+1+diabetes| 11|\n",
      "|BMX-001, a novel ...|[Bruni A, Pepper ...|Am J Transplant. ...|Islet transplanta...|type+1+diabetes| 12|\n",
      "|Continuous Glucos...|[Wood A, O'Neal D...|Intern Med J. 201...|The advent of dev...|type+1+diabetes| 13|\n",
      "|Factitious hypogl...|[Bauman V, Sturke...|Pediatr Diabetes....|BACKGROUND: Facti...|type+1+diabetes| 14|\n",
      "|Intensive remote ...|[Gandrud L, Altan...|Pediatr Diabetes....|OBJECTIVE: While ...|type+1+diabetes| 15|\n",
      "|Lysosomal Exoglyc...|[Maciejczyk M, Ko...|J Diabetes Res. 2...|Before this study...|type+1+diabetes| 16|\n",
      "|MHC-mismatched mi...|[Zhang M, Racine ...|Proc Natl Acad Sc...|Autoimmune type 1...|type+1+diabetes| 17|\n",
      "|Characteristics a...|[Ritsinger V, Her...|Diabetes Care. 20...|OBJECTIVE: To des...|type+1+diabetes| 18|\n",
      "|Parenting Stress ...|[Limbers CA, Teas...|Fam Community Hea...|The purpose of th...|type+1+diabetes| 19|\n",
      "+--------------------+--------------------+--------------------+--------------------+---------------+---+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "non_empty_texts.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_indexer = spark_ft.StringIndexer(inputCol='Topic', outputCol='label')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_indexer_model = label_indexer.fit(non_empty_texts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_deindexer = spark_ft.IndexToString(inputCol='prediction', outputCol='pred_label', \n",
    "                                         labels=label_indexer_model.labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, test = label_indexer_model.transform(non_empty_texts).randomSplit(weights=[0.8, 0.2], seed=123)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "abstract_assembler = DocumentAssembler() \\\n",
    "    .setInputCol(\"Abstract\") \\\n",
    "    .setOutputCol(\"document\")\n",
    "    \n",
    "title_assembler = DocumentAssembler() \\\n",
    "    .setInputCol(\"Title\") \\\n",
    "    .setOutputCol(\"document\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence_detector = SentenceDetector() \\\n",
    "    .setInputCols([\"document\"]) \\\n",
    "    .setOutputCol(\"sentence\") \\\n",
    "    .setUseAbbreviations(True)\n",
    "    \n",
    "tokenizer = Tokenizer() \\\n",
    "  .setInputCols([\"sentence\"]) \\\n",
    "  .setOutputCol(\"token\")\n",
    "\n",
    "stemmer = Stemmer() \\\n",
    "    .setInputCols([\"token\"]) \\\n",
    "    .setOutputCol(\"stem\")\n",
    "    \n",
    "normalizer = Normalizer() \\\n",
    "    .setInputCols([\"stem\"]) \\\n",
    "    .setOutputCol(\"normalized\")\n",
    "\n",
    "nlp_pipeline = Pipeline(stages=[sentence_detector, tokenizer, stemmer, normalizer])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "abstract_finisher = Finisher() \\\n",
    "    .setInputCols([\"normalized\"]) \\\n",
    "    .setOutputCols([\"ntokens\"]) \\\n",
    "    .setOutputAsArray(True) \\\n",
    "    .setCleanAnnotations(True)\n",
    "    \n",
    "title_finisher = Finisher() \\\n",
    "    .setInputCols([\"normalized\"]) \\\n",
    "    .setOutputCols([\"title\"]) \\\n",
    "    .setOutputAsArray(True) \\\n",
    "    .setCleanAnnotations(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "stopWords = spark_ft.StopWordsRemover.loadDefaultStopWords('english')\n",
    "sw_remover = spark_ft.StopWordsRemover(inputCol='ntokens', outputCol='text', stopWords=stopWords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "abstract_pipeline = Pipeline(stages=[abstract_assembler, nlp_pipeline, abstract_finisher, sw_remover])\n",
    "\n",
    "title_pipeline = Pipeline(stages=[title_assembler, nlp_pipeline, title_finisher])\n",
    "\n",
    "preproc_pipeline = Pipeline(stages=[abstract_pipeline, title_pipeline])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "preproc_model = preproc_pipeline.fit(train)\n",
    "processed = preproc_model.transform(train).select('id', 'topic', 'title', 'text', 'label')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+--------------------+--------------------+-----+\n",
      "| id|          topic|               title|                text|label|\n",
      "+---+---------------+--------------------+--------------------+-----+\n",
      "|  0|type+1+diabetes|[a, luciferas, im...|[aim, luciferas, ...|  0.0|\n",
      "|108|type+1+diabetes|[a, typ, diabet, ...|[aimshypothesi, i...|  0.0|\n",
      "|  7|type+1+diabetes|[annal, express, ...|[backgroundto, cl...|  0.0|\n",
      "| 66|type+1+diabetes|[adipos, impact, ...|[object, central,...|  0.0|\n",
      "| 94|type+1+diabetes|[administr, of, v...|[object, two, cas...|  0.0|\n",
      "|122|type+1+diabetes|[alpha, cell, dys...|[typ, diabet, cha...|  0.0|\n",
      "| 77|type+1+diabetes|[an, effect, trea...|[nanotechnologi, ...|  0.0|\n",
      "| 91|type+1+diabetes|[analysi, of, pan...|[background, decr...|  0.0|\n",
      "| 63|type+1+diabetes|[assess, the, nut...|[object, lowcarbo...|  0.0|\n",
      "| 61|type+1+diabetes|[associ, between,...|[aim, investig, a...|  0.0|\n",
      "|  3|type+1+diabetes|[associ, between,...|[aim, depress, mo...|  0.0|\n",
      "|137|type+1+diabetes|[automat, detect,...|[background, auto...|  0.0|\n",
      "| 12|type+1+diabetes|[bmx, a, novel, r...|[islet, transplan...|  0.0|\n",
      "|143|type+1+diabetes|[basal, subnuclea...|[hypothermia, dia...|  0.0|\n",
      "|105|type+1+diabetes|[beta, cell, func...|[background, aim,...|  0.0|\n",
      "| 51|type+1+diabetes|[beta, cell, extr...|[aimshypothesi, i...|  0.0|\n",
      "|116|type+1+diabetes|[bodi, mass, inde...|[object, object, ...|  0.0|\n",
      "|141|type+1+diabetes|[cd, +, t, helper...|[autoreact, cd, +...|  0.0|\n",
      "| 18|type+1+diabetes|[characterist, an...|[object, describ,...|  0.0|\n",
      "| 36|type+1+diabetes|[circul, mirna, p...|[investig, plasma...|  0.0|\n",
      "+---+---------------+--------------------+--------------------+-----+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "processed.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "text2vec = spark_ft.Word2Vec(\n",
    "    vectorSize=100, minCount=5, seed=123, \n",
    "    inputCol='text', outputCol='text_vec', \n",
    "    windowSize=5, maxSentenceLength=30\n",
    ")\n",
    "\n",
    "title2vec = spark_ft.Word2Vec(\n",
    "    vectorSize=50, minCount=3, seed=123, \n",
    "    inputCol='title', outputCol='title_vec', \n",
    "    windowSize=5, maxSentenceLength=10\n",
    ")\n",
    "\n",
    "assembler = spark_ft.VectorAssembler(inputCols=['text_vec', 'title_vec'], outputCol='features')\n",
    "\n",
    "feature_pipeline = Pipeline(stages=[text2vec, title2vec, assembler])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_model = feature_pipeline.fit(processed)\n",
    "features = feature_model.transform(processed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+--------------------+--------------------+-----+--------------------+--------------------+--------------------+\n",
      "| id|          topic|               title|                text|label|            text_vec|           title_vec|            features|\n",
      "+---+---------------+--------------------+--------------------+-----+--------------------+--------------------+--------------------+\n",
      "|  0|type+1+diabetes|[a, luciferas, im...|[aim, luciferas, ...|  0.0|[0.00135010731583...|[-0.0093724036589...|[0.00135010731583...|\n",
      "|108|type+1+diabetes|[a, typ, diabet, ...|[aimshypothesi, i...|  0.0|[-0.0172051572149...|[-0.0485772364307...|[-0.0172051572149...|\n",
      "|  7|type+1+diabetes|[annal, express, ...|[backgroundto, cl...|  0.0|[0.03553665489388...|[-0.0297967683523...|[0.03553665489388...|\n",
      "| 66|type+1+diabetes|[adipos, impact, ...|[object, central,...|  0.0|[0.01813366204482...|[-0.0657454569635...|[0.01813366204482...|\n",
      "| 94|type+1+diabetes|[administr, of, v...|[object, two, cas...|  0.0|[0.00166694424089...|[-0.0769326825315...|[0.00166694424089...|\n",
      "|122|type+1+diabetes|[alpha, cell, dys...|[typ, diabet, cha...|  0.0|[-0.0136437712689...|[-0.0979562591140...|[-0.0136437712689...|\n",
      "| 77|type+1+diabetes|[an, effect, trea...|[nanotechnologi, ...|  0.0|[0.00717100351251...|[-0.0362102882388...|[0.00717100351251...|\n",
      "| 91|type+1+diabetes|[analysi, of, pan...|[background, decr...|  0.0|[0.03914943991343...|[-0.0596927608285...|[0.03914943991343...|\n",
      "| 63|type+1+diabetes|[assess, the, nut...|[object, lowcarbo...|  0.0|[-0.0071556649848...|[-0.0019804228601...|[-0.0071556649848...|\n",
      "| 61|type+1+diabetes|[associ, between,...|[aim, investig, a...|  0.0|[0.04105319693723...|[-0.0671359592815...|[0.04105319693723...|\n",
      "|  3|type+1+diabetes|[associ, between,...|[aim, depress, mo...|  0.0|[0.00475204453800...|[-0.0542618287727...|[0.00475204453800...|\n",
      "|137|type+1+diabetes|[automat, detect,...|[background, auto...|  0.0|[-0.0111896677184...|[-0.0113201303950...|[-0.0111896677184...|\n",
      "| 12|type+1+diabetes|[bmx, a, novel, r...|[islet, transplan...|  0.0|[0.00678624764814...|[-0.0057002936489...|[0.00678624764814...|\n",
      "|143|type+1+diabetes|[basal, subnuclea...|[hypothermia, dia...|  0.0|[-0.0133420593804...|[-0.0176167050696...|[-0.0133420593804...|\n",
      "|105|type+1+diabetes|[beta, cell, func...|[background, aim,...|  0.0|[0.05711436714045...|[-0.0244967469901...|[0.05711436714045...|\n",
      "| 51|type+1+diabetes|[beta, cell, extr...|[aimshypothesi, i...|  0.0|[-0.0042026805076...|[-0.0420467860364...|[-0.0042026805076...|\n",
      "|116|type+1+diabetes|[bodi, mass, inde...|[object, object, ...|  0.0|[0.01275832407116...|[-0.0246792081130...|[0.01275832407116...|\n",
      "|141|type+1+diabetes|[cd, +, t, helper...|[autoreact, cd, +...|  0.0|[-0.0025502139021...|[-0.0012579757429...|[-0.0025502139021...|\n",
      "| 18|type+1+diabetes|[characterist, an...|[object, describ,...|  0.0|[0.05452627044602...|[-0.0370577549465...|[0.05452627044602...|\n",
      "| 36|type+1+diabetes|[circul, mirna, p...|[investig, plasma...|  0.0|[-0.0033266883919...|[-0.0668783794778...|[-0.0033266883919...|\n",
      "+---+---------------+--------------------+--------------------+-----+--------------------+--------------------+--------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "features.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "text2vec_model = text2vec.fit(processed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------+------------------+\n",
      "|       word|        similarity|\n",
      "+-----------+------------------+\n",
      "|       girl|0.8909512758255005|\n",
      "|   leukemia|0.8899913430213928|\n",
      "|       foot|0.8795618414878845|\n",
      "|         dr|0.8758804798126221|\n",
      "|recentonset|0.8757311701774597|\n",
      "|osteoporosi|0.8751940131187439|\n",
      "|     gestat|0.8743801116943359|\n",
      "| neuropathi|0.8684095144271851|\n",
      "|        cgl|0.8683326840400696|\n",
      "|       modi|0.8677959442138672|\n",
      "+-----------+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "text2vec_model.findSynonyms('obes', 10).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------+------------------+\n",
      "|        word|        similarity|\n",
      "+------------+------------------+\n",
      "|      combat| 0.904567301273346|\n",
      "|      experi|0.8929694890975952|\n",
      "|  postdeploy|0.8746237754821777|\n",
      "|  catastroph| 0.864286482334137|\n",
      "|relationship|0.8641922473907471|\n",
      "|         tbi|0.8591167330741882|\n",
      "|         sud|0.8578217029571533|\n",
      "|     symptom|  0.85660320520401|\n",
      "|        ptss|0.8558170199394226|\n",
      "|      buffer|0.8549688458442688|\n",
      "+------------+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "text2vec_model.findSynonyms('trauma', 10).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlpc = spark_cls.MultilayerPerceptronClassifier(\n",
    "    maxIter=100, seed=123, layers=[150, 75, 4]\n",
    ")\n",
    "\n",
    "model_pipeline = Pipeline(stages=[mlpc, label_deindexer])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model_pipeline.fit(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_processed = preproc_model.transform(test).select('id', 'topic', 'title', 'text', 'label')\n",
    "\n",
    "test_features = feature_model.transform(test_processed)\n",
    "\n",
    "preds = model.transform(test_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+--------------------+--------------------+-----+--------------------+--------------------+--------------------+----------+---------------+\n",
      "| id|          topic|               title|                text|label|            text_vec|           title_vec|            features|prediction|     pred_label|\n",
      "+---+---------------+--------------------+--------------------+-----+--------------------+--------------------+--------------------+----------+---------------+\n",
      "|115|type+1+diabetes|[a, physic, activ...|[background, thi,...|  0.0|[-0.0181661611564...|[-0.0456972203346...|[-0.0181661611564...|       0.0|type+1+diabetes|\n",
      "| 47|type+1+diabetes|[a, plateau, in, ...|[object, describ,...|  0.0|[0.02216673717203...|[-0.0611507280264...|[0.02216673717203...|       0.0|type+1+diabetes|\n",
      "| 73|type+1+diabetes|[a, quarter, of, ...|[individu, typ, d...|  0.0|[0.01657626631090...|[-0.0464318374563...|[0.01657626631090...|       0.0|type+1+diabetes|\n",
      "| 69|type+1+diabetes|[allcaus, mortal,...|[object, estim, a...|  0.0|[0.03344887619467...|[-0.0427703198716...|[0.03344887619467...|       0.0|type+1+diabetes|\n",
      "|  4|type+1+diabetes|[alpha, antitryps...|[object, determin...|  0.0|[0.00590007096664...|[-0.0369152471110...|[0.00590007096664...|       0.0|type+1+diabetes|\n",
      "|147|type+1+diabetes|[associ, between,...|[aim, advanc, gly...|  0.0|[0.02942380872589...|[-0.0447189292826...|[0.02942380872589...|       0.0|type+1+diabetes|\n",
      "|127|type+1+diabetes|[cish, promot, po...|[background, impa...|  0.0|[0.02050852430258...|[-0.0283943034716...|[0.02050852430258...|       0.0|type+1+diabetes|\n",
      "| 49|type+1+diabetes|[clinic, profil, ...|[background, diab...|  0.0|[-0.0015870097104...|[-0.0781120069324...|[-0.0015870097104...|       0.0|type+1+diabetes|\n",
      "| 87|type+1+diabetes|[concert, redox, ...|[diabet, cardiome...|  0.0|[-0.0198416023941...|[-0.0224614477949...|[-0.0198416023941...|       0.0|type+1+diabetes|\n",
      "| 90|type+1+diabetes|[construct, of, e...|[pancreat, islet,...|  0.0|[-0.0284984559506...|[-0.0215841557714...|[-0.0284984559506...|       0.0|type+1+diabetes|\n",
      "|114|type+1+diabetes|[differenti, meth...|[diabet, mellitu,...|  0.0|[-0.0276550344460...|[0.01172700314054...|[-0.0276550344460...|       0.0|type+1+diabetes|\n",
      "| 60|type+1+diabetes|[effect, of, resi...|[background, bodi...|  0.0|[0.02043139191001...|[-0.0462926091548...|[0.02043139191001...|       0.0|type+1+diabetes|\n",
      "|107|type+1+diabetes|[evalu, of, vitam...|[background, vita...|  0.0|[0.05376429138559...|[-0.0255385008973...|[0.05376429138559...|       2.0|  heart+disease|\n",
      "| 89|type+1+diabetes|[food, insecur, i...|[aim, household, ...|  0.0|[0.01205506821111...|[-0.0370869886299...|[0.01205506821111...|       0.0|type+1+diabetes|\n",
      "| 37|type+1+diabetes|[genet, of, typ, ...|[typ, diabet, com...|  0.0|[-0.0394111465742...|[-0.1558791380375...|[-0.0394111465742...|       0.0|type+1+diabetes|\n",
      "|152|type+1+diabetes|[hemoglobin, ac, ...|[typ, diabet, td,...|  0.0|[0.03471668307523...|[-0.0538453668954...|[0.03471668307523...|       0.0|type+1+diabetes|\n",
      "| 74|type+1+diabetes|[ildrfc, i, a, no...|[ildr, member, ig...|  0.0|[-0.0393003275928...|[-0.0267670435963...|[-0.0393003275928...|       0.0|type+1+diabetes|\n",
      "| 59|type+1+diabetes|[impact, of, typ,...|[background, neur...|  0.0|[-0.0047581135744...|[-0.0638369887601...|[-0.0047581135744...|       0.0|type+1+diabetes|\n",
      "|109|type+1+diabetes|[impair, hypoglyc...|[hypoglycaemia, r...|  0.0|[-0.0340725541569...|[-0.0500555422157...|[-0.0340725541569...|       0.0|type+1+diabetes|\n",
      "| 40|type+1+diabetes|[influenc, of, hy...|[background, aim,...|  0.0|[0.04033954996631...|[-0.0641401749337...|[0.04033954996631...|       0.0|type+1+diabetes|\n",
      "+---+---------------+--------------------+--------------------+-----+--------------------+--------------------+--------------------+----------+---------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "preds.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_df = preds.select('title', 'text', 'label', 'prediction').toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>text</th>\n",
       "      <th>label</th>\n",
       "      <th>prediction</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[a, physic, activ, intervent, for, children, w...</td>\n",
       "      <td>[background, thi, studi, describ, develop, fea...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[a, plateau, in, new, onset, typ, diabet, inci...</td>\n",
       "      <td>[object, describ, incid, preval, typ, diabet, ...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[a, quarter, of, patient, with, typ, diabet, h...</td>\n",
       "      <td>[individu, typ, diabet, td, ar, increas, risk,...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[allcaus, mortal, in, adult, with, and, withou...</td>\n",
       "      <td>[object, estim, agespecif, sexspecif, allcaus,...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[alpha, antitrypsin, treatment, of, newonset, ...</td>\n",
       "      <td>[object, determin, safeti, pharmacokinet, alph...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               title  \\\n",
       "0  [a, physic, activ, intervent, for, children, w...   \n",
       "1  [a, plateau, in, new, onset, typ, diabet, inci...   \n",
       "2  [a, quarter, of, patient, with, typ, diabet, h...   \n",
       "3  [allcaus, mortal, in, adult, with, and, withou...   \n",
       "4  [alpha, antitrypsin, treatment, of, newonset, ...   \n",
       "\n",
       "                                                text  label  prediction  \n",
       "0  [background, thi, studi, describ, develop, fea...    0.0         0.0  \n",
       "1  [object, describ, incid, preval, typ, diabet, ...    0.0         0.0  \n",
       "2  [individu, typ, diabet, td, ar, increas, risk,...    0.0         0.0  \n",
       "3  [object, estim, agespecif, sexspecif, allcaus,...    0.0         0.0  \n",
       "4  [object, determin, safeti, pharmacokinet, alph...    0.0         0.0  "
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0, 'type+1+diabetes'),\n",
       " (1, 'post+traumatic+stress+disorder'),\n",
       " (2, 'heart+disease'),\n",
       " (3, 'creutzfeldt+jakob+disease')]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(enumerate(label_indexer_model.labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pred type+1+diabetes</th>\n",
       "      <th>pred post+traumatic+stress+disorder</th>\n",
       "      <th>pred heart+disease</th>\n",
       "      <th>pred creutzfeldt+jakob+disease</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>true type+1+diabetes</th>\n",
       "      <td>115</td>\n",
       "      <td>0</td>\n",
       "      <td>9</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>true post+traumatic+stress+disorder</th>\n",
       "      <td>1</td>\n",
       "      <td>98</td>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>true heart+disease</th>\n",
       "      <td>9</td>\n",
       "      <td>1</td>\n",
       "      <td>63</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>true creutzfeldt+jakob+disease</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>92</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                     pred type+1+diabetes  \\\n",
       "true type+1+diabetes                                  115   \n",
       "true post+traumatic+stress+disorder                     1   \n",
       "true heart+disease                                      9   \n",
       "true creutzfeldt+jakob+disease                          0   \n",
       "\n",
       "                                     pred post+traumatic+stress+disorder  \\\n",
       "true type+1+diabetes                                                   0   \n",
       "true post+traumatic+stress+disorder                                   98   \n",
       "true heart+disease                                                     1   \n",
       "true creutzfeldt+jakob+disease                                         0   \n",
       "\n",
       "                                     pred heart+disease  \\\n",
       "true type+1+diabetes                                  9   \n",
       "true post+traumatic+stress+disorder                   5   \n",
       "true heart+disease                                   63   \n",
       "true creutzfeldt+jakob+disease                        3   \n",
       "\n",
       "                                     pred creutzfeldt+jakob+disease  \n",
       "true type+1+diabetes                                              4  \n",
       "true post+traumatic+stress+disorder                               0  \n",
       "true heart+disease                                                1  \n",
       "true creutzfeldt+jakob+disease                                   92  "
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(\n",
    "    data=skmetrics.confusion_matrix(pred_df['label'], pred_df['prediction']),\n",
    "    columns=['pred ' + l for l in label_indexer_model.labels],\n",
    "    index=['true ' + l for l in label_indexer_model.labels]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                precision    recall  f1-score   support\n",
      "\n",
      "               type+1+diabetes       0.92      0.90      0.91       128\n",
      "post+traumatic+stress+disorder       0.99      0.94      0.97       104\n",
      "                 heart+disease       0.79      0.85      0.82        74\n",
      "     creutzfeldt+jakob+disease       0.95      0.97      0.96        95\n",
      "\n",
      "                   avg / total       0.92      0.92      0.92       401\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(skmetrics.classification_report(pred_df['label'], pred_df['prediction'], \n",
    "                                      target_names=label_indexer_model.labels))"
   ]
  }
 ],
 "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}