20 lines
724 B
Python
20 lines
724 B
Python
import unittest
|
|
|
|
from src.sus.lemmatize_text import lemmatize_token
|
|
from src.sus.nlp import nlp
|
|
|
|
|
|
class TestLemmatizeToken(unittest.TestCase):
|
|
def test_simple(self) -> None:
|
|
token = nlp("Center")[0]
|
|
|
|
self.assertEqual(lemmatize_token(token), "centre")
|
|
self.assertEqual(lemmatize_token(token, add_negation=True), "centre")
|
|
self.assertEqual(lemmatize_token(token, add_part_of_speech=True), "centre_NOUN")
|
|
|
|
def test_punctuation(self) -> None:
|
|
token = nlp("This.")[1]
|
|
|
|
self.assertEqual(lemmatize_token(token), ".")
|
|
self.assertEqual(lemmatize_token(token, add_negation=True), ".")
|
|
self.assertEqual(lemmatize_token(token, add_part_of_speech=True), "._PUNCT")
|