Add cli argument for configuring labels

This commit is contained in:
Marianne Stecklina
2019-09-19 11:29:20 +02:00
committed by thomwolf
parent 3e9420add1
commit 99b189df6d
2 changed files with 24 additions and 17 deletions

View File

@@ -202,5 +202,12 @@ def convert_examples_to_features(examples,
return features
def get_labels():
return ["O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]
def get_labels(path):
if path:
with open(path, "r") as f:
labels = f.read().splitlines()
if "O" not in labels:
labels = ["O"] + labels
return labels
else:
return ["O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]