def load_train_test_dfs(local_cache_path="./",
test_fraction=0.5,
random_seed=None):
"""
Get the training and testing data frames based on test_fraction.
Args:
local_cache_path (str): Path to store the data. If the data file
doesn't exist in this path, it's downloaded.
test_fraction (float, optional): Fraction of data ot use for
testing. Since this is a small dataset, the default testing
fraction is set to 0.5
random_seed (float, optional): Random seed used to shuffle the data.
Returns:
tuple: (train_pandas_df, test_pandas_df), each data frame contains
two columns
"sentence": sentences in strings.
"labels": list of entity labels of the words in the sentence.
"""
file_name = URL.split("/")[-1]
maybe_download(URL, file_name, local_cache_path)
data_file = os.path.join(local_cache_path, file_name)
with open(data_file, "r", encoding="utf8") as file:
text = file.read()
sentence_list, labels_list = preprocess_conll(text)
if random_seed:
random.seed(random_seed)
sentence_and_labels = list(zip(sentence_list, labels_list))
random.shuffle(sentence_and_labels)
sentence_list[:], labels_list[:] = zip(*sentence_and_labels)
sentence_count = len(sentence_list)
test_sentence_count = round(sentence_count * test_fraction)
test_sentence_list = sentence_list[:test_sentence_count]
test_labels_list = labels_list[:test_sentence_count]
train_sentence_list = sentence_list[test_sentence_count:]
train_labels_list = labels_list[test_sentence_count:]
train_df = pd.DataFrame({
"sentence": train_sentence_list,
"labels": train_labels_list
})
test_df = pd.DataFrame({
"sentence": test_sentence_list,
"labels": test_labels_list
})
return (train_df, test_df)
def test_ner_utils(ner_utils_test_data):
output = preprocess_conll(ner_utils_test_data["input"])
assert output == ner_utils_test_data["expected_output"]