named-entity-recognition-in-lyrics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook explores using the [Stanford Named Entity Recognizer](https://nlp.stanford.edu/software/CRF-NER.html) and the [Natural Language Tool Kit](http://www.nltk.org/). To run the code you will first need to download the StanfordNER toolkit from [here](https://nlp.stanford.edu/software/CRF-NER.html#Download) and unzip it in the same directory as this notebook file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import datetime \n",
    "import csv\n",
    "from itertools import groupby\n",
    "from nltk.corpus import PlaintextCorpusReader\n",
    "from nltk.tag import StanfordNERTagger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may need to adjust the path `stanford-ner-2018-10-16` if you have downloaded a different version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "tagger = StanfordNERTagger('./stanford-ner-2018-10-16/classifiers/english.all.3class.distsim.crf.ser.gz',\n",
    "                           './stanford-ner-2018-10-16/stanford-ner.jar')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ARTISTS = ['The Roots',\n",
    "'Eve',\n",
    "'DJ Jazzy Jeff & The Fresh Prince',\n",
    "'Ludacris',\n",
    "'T.I.',\n",
    "'Kanye West',\n",
    "'Chance the Rapper',\n",
    "'Common',\n",
    "'Gucci Mane',\n",
    "'Migos',\n",
    "'OutKast',\n",
    "'Twista',\n",
    "'Crucial Conflict',\n",
    "'Lupe Fiasco',\n",
    "'Digital Underground',\n",
    "'2pac',\n",
    "'Trouble Funk']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "LYRICS_DIR = './lyrics/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "paths = [os.path.join(LYRICS_DIR, d) for d in os.listdir(LYRICS_DIR)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tagged_entities(dirpath):\n",
    "    \"\"\"\n",
    "    Given a directory path,\n",
    "    create a PlaintextCorpus object and then\n",
    "    return a list of dictionaries representing tagged entities\n",
    "    in the corpus\n",
    "    \"\"\"\n",
    "    entities = []\n",
    "    corp = PlaintextCorpusReader(dirpath, '.*\\.txt')\n",
    "    \n",
    "    print('{0}: Starting work on {1}'.format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), dirpath))\n",
    "    for fid in corp.fileids():\n",
    "        \n",
    "        s_id = fid.split('-', 1)[0]\n",
    "        s_title = re.sub('-', ' ', fid.split('-', 1)[1].rsplit('.')[0])\n",
    "        \n",
    "        text = corp.words(fid)\n",
    "        tagged = tagger.tag(text)\n",
    "        \n",
    "        # Based on: http://stackoverflow.com/a/30665014/1232820\n",
    "        for tag, chunk in groupby(tagged, lambda x: x[1]):\n",
    "            if tag != 'O':\n",
    "                result = {\n",
    "                    'song_id': s_id,\n",
    "                    'song_title': s_title,\n",
    "                    'entity': \" \".join(word for word, tg in chunk),\n",
    "                    'entity_type': tag\n",
    "                }\n",
    "                entities.append(result)\n",
    "            else:\n",
    "                # flush chunk\n",
    "                pass\n",
    "    print('{0}: Finished processing {1}'.format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), dirpath))\n",
    "    return entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2019-07-17 15:32:50: Starting work on ./lyrics/Twista\n",
      "2019-07-17 15:36:11: Finished processing ./lyrics/Twista\n",
      "2019-07-17 15:36:11: Starting work on ./lyrics/T.I.\n",
      "2019-07-17 15:44:01: Finished processing ./lyrics/T.I.\n",
      "2019-07-17 15:44:01: Starting work on ./lyrics/The Roots\n",
      "2019-07-17 15:50:05: Finished processing ./lyrics/The Roots\n",
      "2019-07-17 15:50:05: Starting work on ./lyrics/2Pac\n",
      "2019-07-17 16:04:08: Finished processing ./lyrics/2Pac\n",
      "2019-07-17 16:04:08: Starting work on ./lyrics/Digital Underground\n",
      "2019-07-17 16:05:06: Finished processing ./lyrics/Digital Underground\n",
      "2019-07-17 16:05:06: Starting work on ./lyrics/DJ Jazzy Jeff & The Fresh Prince\n",
      "2019-07-17 16:06:28: Finished processing ./lyrics/DJ Jazzy Jeff & The Fresh Prince\n",
      "2019-07-17 16:06:28: Starting work on ./lyrics/Crucial Conflict\n",
      "2019-07-17 16:06:59: Finished processing ./lyrics/Crucial Conflict\n",
      "2019-07-17 16:06:59: Starting work on ./lyrics/Common\n",
      "2019-07-17 16:13:46: Finished processing ./lyrics/Common\n",
      "2019-07-17 16:13:46: Starting work on ./lyrics/Trouble Funk\n",
      "2019-07-17 16:13:50: Finished processing ./lyrics/Trouble Funk\n",
      "2019-07-17 16:13:50: Starting work on ./lyrics/Eve\n",
      "2019-07-17 16:14:57: Finished processing ./lyrics/Eve\n",
      "2019-07-17 16:14:57: Starting work on ./lyrics/Migos\n",
      "2019-07-17 16:24:45: Finished processing ./lyrics/Migos\n",
      "2019-07-17 16:24:45: Starting work on ./lyrics/Lupe Fiasco\n",
      "2019-07-17 16:36:18: Finished processing ./lyrics/Lupe Fiasco\n",
      "2019-07-17 16:36:18: Starting work on ./lyrics/Gucci Mane\n",
      "2019-07-17 17:00:26: Finished processing ./lyrics/Gucci Mane\n",
      "2019-07-17 17:00:26: Starting work on ./lyrics/Kanye West\n",
      "2019-07-17 17:13:51: Finished processing ./lyrics/Kanye West\n",
      "2019-07-17 17:13:51: Starting work on ./lyrics/OutKast\n",
      "2019-07-17 17:18:30: Finished processing ./lyrics/OutKast\n",
      "2019-07-17 17:18:30: Starting work on ./lyrics/Chance The Rapper\n",
      "2019-07-17 17:22:05: Finished processing ./lyrics/Chance The Rapper\n",
      "2019-07-17 17:22:05: Starting work on ./lyrics/Ludacris\n",
      "2019-07-17 17:28:20: Finished processing ./lyrics/Ludacris\n"
     ]
    }
   ],
   "source": [
    "all_entities = [get_tagged_entities(p) for p in paths]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "merge = zip(sorted(ARTISTS), all_entities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "for m in merge:\n",
    "    update = {'artist': m[0]}\n",
    "    new_item = [{**item, **update} for item in m[1]]\n",
    "    if len(m[1]) == len(new_item):\n",
    "        results.extend(new_item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "15819"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./entities.csv', 'w') as csvfile:\n",
    "    field_names = ['entity', 'entity_type', 'song_id', 'song_title', 'artist']\n",
    "    writer = csv.DictWriter(csvfile, fieldnames=field_names)\n",
    "    \n",
    "    writer.writeheader()\n",
    "    for row in results:\n",
    "        writer.writerow(row)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}