-
Notifications
You must be signed in to change notification settings - Fork 0
/
load_ner_2014.py
69 lines (50 loc) · 2.05 KB
/
load_ner_2014.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import transformers as ts
import datasets as ds
# SET DIRS
# Set directories to wherever the data is stored after preprocessing based on:
raw_data_dir = "/mnt/sdd/niallt/bio-lm/data/tasks/I2B22014NER/" #CHANGEME
save_data_dir = "/mnt/sdd/niallt/bio-lm/data/tasks/i2b2-2014_hf_dataset/" #CHANGEME
allLabels = []
dataDict = {}
def load_ner_dataset(path, subset):
dataDict[subset] = {
"tokens": [],
"ner_tags_str": [],
"ner_tags": [],
}
lines = []
with open(path, mode="r") as f:
lines = f.readlines()
sentences = []
labels = []
currentSampleTokens = []
currentSampleLabels = []
for line in lines:
if line.strip() == "":
sentences.append(currentSampleTokens)
labels.append(currentSampleLabels)
currentSampleTokens = []
currentSampleLabels = []
else:
cleanedLine = line.replace("\n", "")
token, label = cleanedLine.split(
" ")[0].strip(), cleanedLine.split(" ")[1].strip()
currentSampleTokens.append(token)
currentSampleLabels.append(label)
allLabels.append(label)
dataDict[subset]["tokens"] = sentences
dataDict[subset]["ner_tags_str"] = labels
# Set directories to wherever the data is stored after preprocessing based on: https://github.com/facebookresearch/bio-lm/tree/main/preprocessing
load_ner_dataset(f"{raw_data_dir}/train.txt.conll", "train")
load_ner_dataset(f"{raw_data_dir}/dev.txt.conll", "validation")
load_ner_dataset(f"{raw_data_dir}/test.txt.conll", "test")
allLabels = list(set(allLabels))
label_to_index = {label: index for index, label in enumerate(allLabels)}
for key, value in dataDict.items():
dataDict[key]["ner_tags"] = [[label_to_index[label] for label in str_labels] for str_labels in value["ner_tags_str"]]
dataDict[key] = ds.Dataset.from_dict(dataDict[key])
dataDict["info"] = ds.Dataset.from_dict({"all_ner_tags": [allLabels]})
dataset = ds.DatasetDict(dataDict)
print(dataset)
# Set location to save the dataset to - this will be a huggingface dataset
dataset.save_to_disk(f"{save_data_dir}")