Named Entity Chunker

.travis.yml

@@ -1,4 +1,4 @@
-sudo: false
+sudo: true
 language: python
 
 python:
@@ -7,9 +7,13 @@ python:
   - '2.7'
   - pypy
 
+addons:
+  apt:
+    packages:
+    - numpy
+
 install:
   - pip install coveralls
-  - pip install nltk
   - pip install -r requirements.txt
 
 services:

chatterbot/utils/entity_tagger.py

@@ -0,0 +1,48 @@
+import nltk
+
+from chatterbot.utils.pos_tagger import POSTagger
+
+
+class NamedEntityTagger():
+    """
+    A wrapper class for ChatterBot's custom implementations
+    of:
+    1) Named Entity chunking
+    """
+
+    def __init__(self):
+        """
+        Constructor to initialize instance variables.
+        """
+        from nltk.data import find
+        from nltk import download
+
+        try:
+            find('words.zip')
+        except LookupError:
+            download('words')
+
+        try:
+            find('maxent_ne_chunker.zip')
+        except LookupError:
+            download('maxent_ne_chunker')
+
+        self.tagger = POSTagger()
+
+    def ne_chunk(self, string):
+        """
+        Find all of the named entities and return them.
+        """
+        ne_list = []
+
+        named_entities = nltk.ne_chunk(self.tagger.tag(self.tagger.tokenize(string)), binary=True)
+        named_entities = nltk.chunk.tree2conlltags(named_entities)
+
+        # Getting named entities in a text
+        for entity in named_entities:
+            if "NE" in entity[2]:
+                if entity not in ne_list:
+                    ne_list.append(entity[0])
+
+        # Returning list of named entities
+        return ne_list

chatterbot/utils/pos_tagger.py

@@ -1,7 +1,9 @@
 from nltk import word_tokenize
+from nltk import pos_tag
 
 
 class POSTagger():
+
     def __init__(self):
         from nltk.data import find
         from nltk import download
@@ -11,6 +13,11 @@ def __init__(self):
         except LookupError:
             download('punkt')
 
+        try:
+            find('averaged_perceptron_tagger.zip')
+        except LookupError:
+            download('averaged_perceptron_tagger')
+
     def tokenize(self, text):
         """
         Takes an input string and tokenizes that text.
@@ -19,3 +26,10 @@ def tokenize(self, text):
         """
 
         return word_tokenize(text)
+
+    def tag(self, tokens):
+        """
+        Takes a set of tokens and returns the tagged tokens.
+        """
+
+        return pos_tag(tokens)

tests/test_utils.py

@@ -7,6 +7,7 @@
 from chatterbot.utils.pos_tagger import POSTagger
 from chatterbot.utils.stop_words import StopWordsManager
 from chatterbot.utils.word_net import Wordnet
+from chatterbot.utils.entity_tagger import NamedEntityTagger
 
 
 class UtilityTests(TestCase):
@@ -34,6 +35,14 @@ def test_word_net(self):
 
         self.assertEqual(0.06666666666666667, synsets[0].path_similarity(synsets[1]))
 
+    def test_ne_tagger(self):
+        named_entity_tagger = NamedEntityTagger()
+        named_entities = named_entity_tagger.ne_chunk("test one, two, and three.")
+        valid_check = named_entity_tagger.ne_chunk("New York Times")
+
+        self.assertEqual(named_entities, [])
+        self.assertEqual(valid_check, ['New', 'York', 'Times'])
+
 class CleanWhitespaceTests(TestCase):
 
     def test_clean_whitespace(self):