create-test-data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4bcff7ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import spacy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "93472cf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_data(file, data):\n",
    "    with open (file, \"w\", encoding=\"utf-8\") as f:\n",
    "        json.dump(data, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "242b2501",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"commodities_ner_rules\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "63bf9122",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_results(text):\n",
    "    doc = nlp(text)\n",
    "    results = []\n",
    "    entities = []\n",
    "    for ent in doc.ents:\n",
    "        entities.append((ent.start_char, ent.end_char, ent.label_))\n",
    "    if len(entities) > 0:\n",
    "        results = [text, {\"entities\": entities}]\n",
    "        return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b7556e64",
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST_DATA = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1f9c1282",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['The sanctions, which were announced on April 6, forced companies around the world to stop doing business with Rusal, the biggest aluminium producer outside of China. But under the revised guidance from the US Treasury on Monday, Rusal’s non-US counterparties now have until October to wind up business with the company and can make payments to non-blocked bank accounts.', {'entities': [(129, 138, 'COMMODITY')]}]\n",
      "2798\n"
     ]
    }
   ],
   "source": [
    "with open(\"data/ft-articles-test.txt\") as file:\n",
    "    text = file.read()\n",
    "    articles = text.split(\"\\n\")\n",
    "    for article in articles:\n",
    "        articleUuid, articleBodyText = article.split(\"|||\")\n",
    "        segments = articleBodyText.split(\"||\")\n",
    "        hits = []\n",
    "        for segment in segments:\n",
    "            results = extract_results(segment)\n",
    "            if results != None:\n",
    "                TEST_DATA.append(results)\n",
    "                \n",
    "print(TEST_DATA[0])\n",
    "print(len(TEST_DATA))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "af3fbdfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_data(\"data/test_data.json\", TEST_DATA)"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}