def example_5_1():
happy_toc = HappyTokenClassification(model_type="BERT",
model_name="dslim/bert-base-NER")
result = happy_toc.classify_token(
"My name is Geoffrey and I live in Toronto")
print(type(result)) # <class 'list'>
print(result[0].word) # Geoffrey
print(result[0].entity) # B-PER
print(result[0].score) # 0.9988969564437866
print(result[0].index) # 4
print(result[0].start) # 11
print(result[0].end) # 19
print(result[1].word) # Toronto
print(result[1].entity) # B-LOC
def test_toc_save():
happy = HappyTokenClassification()
happy.save("model/")
result_before = happy.classify_token(
"My name is Geoffrey and I live in Toronto")
happy = HappyTokenClassification(load_path="model/")
result_after = happy.classify_token(
"My name is Geoffrey and I live in Toronto")
assert result_before[0].word == result_after[0].word
def test_classify_text():
happy_toc = HappyTokenClassification(model_type="BERT",
model_name="dslim/bert-base-NER")
expected_result = [
TokenClassificationResult(word='Geoffrey',
score=0.9988969564437866,
entity='B-PER',
index=4,
start=11,
end=19),
TokenClassificationResult(word='Toronto',
score=0.9993201494216919,
entity='B-LOC',
index=9,
start=34,
end=41)
]
result = happy_toc.classify_token(
"My name is Geoffrey and I live in Toronto")
assert result == expected_result
def test_classify_text():
happy_toc = HappyTokenClassification(model_type="BERT",
model_name="dslim/bert-base-NER")
expected_result = [{
'word': 'Geoffrey',
'score': 0.9988969564437866,
'entity': 'B-PER',
'index': 4,
'start': 11,
'end': 19
}, {
'word': 'Toronto',
'score': 0.9993201494216919,
'entity': 'B-LOC',
'index': 9,
'start': 34,
'end': 41
}]
result = happy_toc.classify_token(
"My name is Geoffrey and I live in Toronto")
assert result == expected_result
def example_5_0():
happy_toc = HappyTokenClassification("BERT",
"dslim/bert-base-NER") # default
happy_toc_large = HappyTokenClassification(
"XLM-ROBERTA", "xlm-roberta-large-finetuned-conll03-english")