pkgs/development/python-modules/spacy-transformers/annotation-test/annotate.py

   1 import pytest
   2 import spacy
   3
   4 en_text = (
   5     "When Sebastian Thrun started working on self-driving cars at "
   6     "Google in 2007, few people outside of the company took him "
   7     "seriously. “I can tell you very senior CEOs of major American "
   8     "car companies would shake my hand and turn away because I wasn’t "
   9     "worth talking to,” said Thrun, in an interview with Recode earlier "
  10     "this week.")
  11
  12
  13 @pytest.fixture
  14 def en_core_web_trf():
  15     return spacy.load("en_core_web_trf")
  16
  17
  18 @pytest.fixture
  19 def doc_en_core_web_trf(en_core_web_trf):
  20     return en_core_web_trf(en_text)
  21
  22
  23 def test_entities(doc_en_core_web_trf):
  24     entities = list(map(lambda e: (e.text, e.label_),
  25                         doc_en_core_web_trf.ents))
  26
  27     assert entities == [
  28         ('Sebastian Thrun', 'PERSON'),
  29         ('Google', 'ORG'),
  30         ('2007', 'DATE'),
  31         ('American', 'NORP'),
  32         ('Thrun', 'PERSON'),
  33         ('Recode', 'ORG'),
  34         ('earlier this week', 'DATE'),
  35     ]
  36
  37
  38 def test_nouns(doc_en_core_web_trf):
  39     assert [
  40         chunk.text for chunk in doc_en_core_web_trf.noun_chunks] == [
  41         'Sebastian Thrun',
  42         'self-driving cars',
  43         'Google',
  44         'few people',
  45         'the company',
  46         'him',
  47         'I',
  48         'you',
  49         'very senior CEOs',
  50         'major American car companies',
  51         'my hand',
  52         'I',
  53         'Thrun',
  54         'an interview',
  55         'Recode']
  56
  57
  58 def test_verbs(doc_en_core_web_trf):
  59     assert [
  60         token.lemma_ for token in doc_en_core_web_trf if token.pos_ == "VERB"] == [
  61         'start',
  62         'work',
  63         'drive',
  64         'take',
  65         'tell',
  66         'shake',
  67         'turn',
  68         'talk',
  69         'say']