1 - Filter Data & Extract Annotations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "acfb0115",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "import time\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import requests\n",
    "from typing import Any\n",
    "from bs4 import BeautifulSoup\n",
    "import functools \n",
    "import re\n",
    "\n",
    "pd.set_option('display.max_rows', 2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8cdf02a",
   "metadata": {},
   "source": [
    "The original dataset is taken from [Kaggle Genius Song Lyrics Dataset](https://www.kaggle.com/datasets/carlosgdcj/genius-song-lyrics-with-language-information?resource=download). It contains list of musical lyrics and poetry along with metadata. The lyrics and ids are later used to pull annotations from Genius and merge them with the lyrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ee20867",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('./data/poets.csv') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "da346159",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_english_language_poets(file_name: str = \"./../data/List of English-language poets - Wikipedia.html\"):   \n",
    "    \"\"\"\n",
    "    Function for extraction of english poets names from Wikipedia\n",
    "    Args:\n",
    "    file_name - HTML file with poets page\n",
    "    Returns:\n",
    "    names_list - list of names\n",
    "    \"\"\"\n",
    "    with open(file_name) as fp:\n",
    "        soup = BeautifulSoup(fp, \"html.parser\")\n",
    "        \n",
    "    all_containers = soup.find_all(\"div\", class_=\"div-col\")\n",
    "    name_elements = []\n",
    "    for c in all_containers:\n",
    "        name_elements += c.findAll('li')\n",
    "    names_list = map(lambda n: n.get_text(), name_elements)\n",
    "    names_list = set(map(lambda n: re.sub(\"[\\(\\[].*?[\\)\\]]\", \"\", n)[:-1], names_list))\n",
    "    \n",
    "    return names_list         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0f65df5",
   "metadata": {},
   "outputs": [],
   "source": [
    "poets = get_english_language_poets()\n",
    "df = df[df['artist'].isin(poets)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea1d0428",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['lyrics_paragraphs_count'] = df['lyrics'].apply(lambda l: len(l.split('\\n\\n')))\n",
    "df['len_words'] = df['lyrics'].apply(lambda l: len(l.split(' ')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9d1aa08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "artists_to_filter = [\n",
    "    'Kendrick Lamar',\n",
    "    'Bob Dylan',\n",
    "    'Christopher Nolan',\n",
    "    'Henry Fielding',\n",
    "    'Nick Cave',\n",
    "    'Charlotte Lennox',\n",
    "    'Ryan Adams'\n",
    "]\n",
    "\n",
    "df = df[(~df['artist'].isin(artists_to_filter))\n",
    "        & (~df['lyrics'].str.contains(\"Verse|Hook|Intro\"))\n",
    "        & (df['language'] == \"en\")\n",
    "        & (~df['title'].str.contains(\"Act |ACT |Chapter |Chap. |CHAPTER |Ch. |CHAP.\"))\n",
    "        & (df['len_words'] <= 650)\n",
    "       ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "25b29c15",
   "metadata": {},
   "outputs": [],
   "source": [
    "ids = sorted(list(map(lambda i: int(i), df['id'].astype(int).values)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d38f149f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19203"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ids) # number of unique song ids"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8388a34",
   "metadata": {},
   "source": [
    "### Get Annotataions from Genius"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9e41825e",
   "metadata": {},
   "outputs": [],
   "source": [
    "access_token # access token from Genius API "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77f4a543",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_referents(ids: list) -> None:\n",
    "    \"\"\"\n",
    "    Method to download referents and annotations for every song id\n",
    "    \"\"\"\n",
    "    for i in range(len(ids)):\n",
    "        counter = i\n",
    "        song_id = ids[counter]\n",
    "        r = requests.get(f'https://api.genius.com/referents?song_id={song_id}&access_token={access_token}')\n",
    "        if r.status_code == 200:\n",
    "            with open(f'./data/referents/{song_id}.json', 'w') as f:\n",
    "                json.dump(r.json(), f)\n",
    "            if counter%20 == 0:\n",
    "                time.sleep(1)\n",
    "                print(f'processed {counter} poems')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "33be2c20",
   "metadata": {},
   "outputs": [],
   "source": [
    "def traverse_annotation_children(el: Any, annotations: list):\n",
    "    \"\"\"\n",
    "    Recursive retrieval of annotation children of the referents. Merges them into 1D list.\n",
    "    \"\"\"\n",
    "    if type(el) == dict:\n",
    "        return traverse_annotation_children(el['children'], annotations)\n",
    "    for c in el:\n",
    "        if (type(c) == list and len(c) == 1) or type(c) == str:\n",
    "            annotations.append(c)\n",
    "        else:\n",
    "            if 'children' in c:\n",
    "                traverse_annotation_children(c['children'], annotations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d3568f68",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_structured_annotation(referents_object: list):\n",
    "    \"\"\"\n",
    "    Flattens the annotations data structure into the list of annotations\n",
    "    \"\"\"\n",
    "    outputs = []\n",
    "    for ref in referents_object:\n",
    "        output = {\n",
    "            'fragment': ref['fragment'],\n",
    "            'classification': ref['classification'],\n",
    "            'annotations': []\n",
    "        }\n",
    "        annotations = ref['annotations']\n",
    "        for a in annotations:\n",
    "            annotation_object = {\n",
    "                'verified': a['verified']\n",
    "            }\n",
    "            if 'children' in a['body']['dom']:\n",
    "                children = a['body']['dom']['children']\n",
    "                annotated_children = []\n",
    "                traverse_annotation_children(children, annotated_children)\n",
    "                annotated_children = list(map(lambda t: t if len(t) > 0 else '\\n', annotated_children))\n",
    "                annotated_children = list(filter(lambda t: not 'http' in t, annotated_children))\n",
    "                if len(annotated_children):\n",
    "                    if len(annotated_children[-1]) == 0:\n",
    "                        annotated_children = annotated_children[:len(annotated_children)-2]\n",
    "                    annotated_children = ' '.join(annotated_children)\n",
    "                    annotation_object['text'] = annotated_children\n",
    "                    output['annotations'].append(annotation_object)\n",
    "        if len(output['annotations']):\n",
    "            outputs.append(output)\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e60a38de",
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_lyrics_with_annotation_by_paragraph(annotation_object: dict, lyrics: str):\n",
    "    \"\"\"\n",
    "    # Deprecated #\n",
    "    Function that adds context to each referent with annotion.\n",
    "    Creates context based on paragraph, not the lines.\n",
    "    1. Splits into paragraphs '\\n\\n'\n",
    "    2. Finds the paragraph of the referent.\n",
    "    3. Does the same logic as \"merge_lyrics_with_annotation\" but only within the paragraph.\n",
    "    \"\"\"\n",
    "    LINES_BEFORE = 8\n",
    "    LINES_AFTER = 8\n",
    "    \n",
    "    paragraphs = list(filter(lambda l: len(l) > 0, lyrics.replace(\"’\", \"'\").split('\\n\\n')))\n",
    "    avg_lines_in_paragraph = np.mean([len(p.split('\\n')) for p in paragraphs])\n",
    "    paragraphs_lines = dict(zip(paragraphs, [p.replace(\"’\", \"'\").split('\\n') for p in paragraphs]))\n",
    "    #poem_lines = list(filter(lambda l: len(l) > 0, lyrics.replace(\"’\", \"'\").split('\\n')))\n",
    "    for fragment in annotation_object:\n",
    "        referent = fragment['fragment'].replace(\"’\", \"'\")\n",
    "        if '\\n' in referent:\n",
    "            if referent[-1] == '\\n':\n",
    "                referent = referent[:-1]\n",
    "            referent_list = referent.split('\\n')\n",
    "            referent_start, referent_end = referent_list[0], referent_list[-1]\n",
    "        else:\n",
    "            referent_start, referent_end = referent, referent\n",
    "        \n",
    "        paragraph_start = list(filter(lambda p: referent_start in p, paragraphs))[0]\n",
    "        paragraph_end = list(filter(lambda p: referent_end in p, paragraphs))[0]\n",
    "        \n",
    "        index_start = paragraphs_lines[paragraph_start].index(referent_start)\n",
    "        index_end = paragraphs_lines[paragraph_end].index(referent_end)\n",
    "                \n",
    "        context_before = paragraphs_lines[paragraph_start][max(index_start-LINES_BEFORE, 0):index_start]\n",
    "        context_after = paragraphs_lines[paragraph_end][index_end+1:min(index_end+LINES_AFTER, len(paragraphs_lines[paragraph_end]))]\n",
    "        print(context_before)\n",
    "        print(referent)\n",
    "        print(context_after)\n",
    "        print('\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "7211a86f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_lyrics_with_annotation(annotation_object: dict, metadata: dict):\n",
    "    \"\"\"\n",
    "    Function that adds context to each referent with annotion.\n",
    "    Currently returns result if referent is fully present in the lyrics (full match).\n",
    "    \"\"\"\n",
    "    LINES_BEFORE = 12\n",
    "    LINES_AFTER = 4\n",
    "    escape_chars_filter = ''.join([chr(i) for i in range(1, 32)]).replace('\\n','') + '\\xa0'\n",
    "    lyrics = metadata['lyrics']\n",
    "    \n",
    "    poem_lines = list(filter(lambda l: len(l) > 0, lyrics.replace(\"’\", \"'\").split('\\n')))\n",
    "    poem_lines = list(map(lambda l: l.translate(str.maketrans('', '', escape_chars_filter)), poem_lines))\n",
    "    results = []\n",
    "    \n",
    "    for fragment in annotation_object:\n",
    "        referent = fragment['fragment'].replace(\"’\", \"'\").translate(str.maketrans('', '', escape_chars_filter)) \n",
    "        if '\\n' in referent:\n",
    "            if referent[-1] == '\\n':\n",
    "                referent = referent[:-1]\n",
    "            referent_list = referent.split('\\n')\n",
    "            referent_start, referent_end = referent_list[0], referent_list[-1]\n",
    "        else:\n",
    "            referent_start, referent_end = referent, referent\n",
    "            \n",
    "        index_start, index_end = -1, -1\n",
    "        \n",
    "#         if referent_start not in poem_lines \\\n",
    "#             and referent_end not in poem_lines:\n",
    "#             #and referent_start == referent_end:\n",
    "#             print('*')\n",
    "#             target_poem_lines = list(filter(lambda p: referent_start.lower() in p.lower(), poem_lines))\n",
    "#             if len(target_poem_lines):\n",
    "#                 index_start = poem_lines.index(target_poem_lines[0])\n",
    "#                 index_end = poem_lines.index(target_poem_lines[0])\n",
    "\n",
    "        if referent_start in poem_lines and referent_end in poem_lines:\n",
    "            index_start = poem_lines.index(referent_start)\n",
    "            index_end = poem_lines.index(referent_end)\n",
    "            \n",
    "        if index_start > -1 and index_end > -1:\n",
    "            context_before = poem_lines[max(index_start-LINES_BEFORE, 0):index_start]\n",
    "            context_after = poem_lines[index_end+1:min(index_end+LINES_AFTER, len(poem_lines))]\n",
    "            results.append({\n",
    "                \"content_before\": '\\n'.join(context_before),\n",
    "                \"referent\": referent,\n",
    "                \"context_after\": '\\n'.join(context_after),\n",
    "                \"annotation\": fragment['annotations'][0]['text'],\n",
    "                \"poet\": metadata['artist'],\n",
    "                \"poem_title\": metadata['title']\n",
    "            })\n",
    "            \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "16dbf8b6",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def json_to_annotation_dicts(folder_path: str = './data/referents', n: int = -1) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Opens json files with referents response from Genius and structures them into the dataset with annotations.\n",
    "    \"\"\"\n",
    "    files = os.listdir(folder_path)\n",
    "    files.remove('.DS_Store')\n",
    "    output_df = []\n",
    "    if n != -1:\n",
    "        files = files[:n]\n",
    "    for file_name in files:\n",
    "        poem_id = int(file_name.replace('.json', ''))\n",
    "        with open(f\"{folder_path}/{file_name}\", \"r\") as f:\n",
    "            referents = json.load(f)['response']['referents']\n",
    "            if len(referents):\n",
    "                referents = list(filter(lambda r: r['classification'] == 'accepted', get_structured_annotation(referents)))\n",
    "            if len(referents):\n",
    "                output_df += merge_lyrics_with_annotation(referents, df[df['id'] == poem_id][['lyrics', 'title', 'artist']].to_dict(orient='records')[0])\n",
    "    return pd.DataFrame(output_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "64ea91e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotations_df = json_to_annotation_dicts()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8cc2ef6",
   "metadata": {},
   "source": [
    "#### Size and sample of the constructed dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "7529bde4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3629, 1479, 280)"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(annotations_df), annotations_df['poem_title'].nunique(), annotations_df['poet'].nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "e690e5e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>content_before</th>\n",
       "      <th>referent</th>\n",
       "      <th>context_after</th>\n",
       "      <th>annotation</th>\n",
       "      <th>poet</th>\n",
       "      <th>poem_title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>351</th>\n",
       "      <td>Some say the world will end in fire,\\nSome say...</td>\n",
       "      <td>But if it had to perish twice,</td>\n",
       "      <td>I think I know enough of hate\\nTo say that for...</td>\n",
       "      <td>In the previous line, Frost claims that he fav...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Fire and Ice</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>352</th>\n",
       "      <td>Some say the world will end in fire,\\nSome say...</td>\n",
       "      <td>And would suffice.</td>\n",
       "      <td></td>\n",
       "      <td>\\n Frost’s ironic understatement reminds us th...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Fire and Ice</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>353</th>\n",
       "      <td>Some say the world will end in fire,\\nSome say...</td>\n",
       "      <td>I think I know enough of hate\\nTo say that for...</td>\n",
       "      <td>And would suffice.</td>\n",
       "      <td>\\n When someone’s words have an “icy” tone, th...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Fire and Ice</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>354</th>\n",
       "      <td>Some say the world will end in fire,\\nSome say...</td>\n",
       "      <td>From what I've tasted of desire\\nI hold with t...</td>\n",
       "      <td>But if it had to perish twice,\\nI think I know...</td>\n",
       "      <td>“Fire and Ice” was written in the aftermath of...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Fire and Ice</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>355</th>\n",
       "      <td></td>\n",
       "      <td>Some say the world will end in fire,\\nSome say...</td>\n",
       "      <td>From what I've tasted of desire\\nI hold with t...</td>\n",
       "      <td>Frost may be alluding to Dante’s  Inferno  (in...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Fire and Ice</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>359</th>\n",
       "      <td>And makes gaps even two can pass abreast.\\nThe...</td>\n",
       "      <td>To each the boulders that have fallen to each.</td>\n",
       "      <td>And some are loaves and some so nearly balls\\n...</td>\n",
       "      <td>Though they are working at the same time, toge...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Mending Wall</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>360</th>\n",
       "      <td>Something there is that doesn't love a wall,\\n...</td>\n",
       "      <td>But at spring mending-time we find them there.</td>\n",
       "      <td>I let my neighbor know beyond the hill;\\nAnd o...</td>\n",
       "      <td>Compare the Greek legend of  Sisyphus,  condem...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Mending Wall</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>361</th>\n",
       "      <td>A Christmas Circular Letter\\nThe city had with...</td>\n",
       "      <td>A-buttoning coats to ask him who he was.</td>\n",
       "      <td>He proved to be the city come again\\nTo look f...</td>\n",
       "      <td>People have to button their coats so they can ...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Christmas Trees</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>589</th>\n",
       "      <td>Here come real stars to fill the upper skies,\\...</td>\n",
       "      <td>Only, of course, they can't sustain the part.</td>\n",
       "      <td></td>\n",
       "      <td>Nothing Gold Can Stay \\n People inevitably yie...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Fireflies in the Garden</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>590</th>\n",
       "      <td>Here come real stars to fill the upper skies,\\...</td>\n",
       "      <td>(And they were never really stars at heart)</td>\n",
       "      <td>Achieve at times a very star-like start.\\nOnly...</td>\n",
       "      <td>In Reality… \\n Fireflies are just bugs. \\n \\n ...</td>\n",
       "      <td>Robert Frost</td>\n",
       "      <td>Fireflies in the Garden</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                        content_before  \\\n",
       "351  Some say the world will end in fire,\\nSome say...   \n",
       "352  Some say the world will end in fire,\\nSome say...   \n",
       "353  Some say the world will end in fire,\\nSome say...   \n",
       "354  Some say the world will end in fire,\\nSome say...   \n",
       "355                                                      \n",
       "359  And makes gaps even two can pass abreast.\\nThe...   \n",
       "360  Something there is that doesn't love a wall,\\n...   \n",
       "361  A Christmas Circular Letter\\nThe city had with...   \n",
       "589  Here come real stars to fill the upper skies,\\...   \n",
       "590  Here come real stars to fill the upper skies,\\...   \n",
       "\n",
       "                                              referent  \\\n",
       "351                     But if it had to perish twice,   \n",
       "352                                 And would suffice.   \n",
       "353  I think I know enough of hate\\nTo say that for...   \n",
       "354  From what I've tasted of desire\\nI hold with t...   \n",
       "355  Some say the world will end in fire,\\nSome say...   \n",
       "359     To each the boulders that have fallen to each.   \n",
       "360     But at spring mending-time we find them there.   \n",
       "361           A-buttoning coats to ask him who he was.   \n",
       "589      Only, of course, they can't sustain the part.   \n",
       "590        (And they were never really stars at heart)   \n",
       "\n",
       "                                         context_after  \\\n",
       "351  I think I know enough of hate\\nTo say that for...   \n",
       "352                                                      \n",
       "353                                 And would suffice.   \n",
       "354  But if it had to perish twice,\\nI think I know...   \n",
       "355  From what I've tasted of desire\\nI hold with t...   \n",
       "359  And some are loaves and some so nearly balls\\n...   \n",
       "360  I let my neighbor know beyond the hill;\\nAnd o...   \n",
       "361  He proved to be the city come again\\nTo look f...   \n",
       "589                                                      \n",
       "590  Achieve at times a very star-like start.\\nOnly...   \n",
       "\n",
       "                                            annotation          poet  \\\n",
       "351  In the previous line, Frost claims that he fav...  Robert Frost   \n",
       "352  \\n Frost’s ironic understatement reminds us th...  Robert Frost   \n",
       "353  \\n When someone’s words have an “icy” tone, th...  Robert Frost   \n",
       "354  “Fire and Ice” was written in the aftermath of...  Robert Frost   \n",
       "355  Frost may be alluding to Dante’s  Inferno  (in...  Robert Frost   \n",
       "359  Though they are working at the same time, toge...  Robert Frost   \n",
       "360  Compare the Greek legend of  Sisyphus,  condem...  Robert Frost   \n",
       "361  People have to button their coats so they can ...  Robert Frost   \n",
       "589  Nothing Gold Can Stay \\n People inevitably yie...  Robert Frost   \n",
       "590  In Reality… \\n Fireflies are just bugs. \\n \\n ...  Robert Frost   \n",
       "\n",
       "                  poem_title  \n",
       "351             Fire and Ice  \n",
       "352             Fire and Ice  \n",
       "353             Fire and Ice  \n",
       "354             Fire and Ice  \n",
       "355             Fire and Ice  \n",
       "359             Mending Wall  \n",
       "360             Mending Wall  \n",
       "361          Christmas Trees  \n",
       "589  Fireflies in the Garden  \n",
       "590  Fireflies in the Garden  "
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "annotations_df[annotations_df['poet'] == 'Robert Frost'].head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "2dec8d6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotations_df.to_csv('./data/annotations_dataset.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "custom_llm",
   "language": "python",
   "name": "custom_llm"
  },
  "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}