# ---------------------------------------------------------------------------- #
# Imports #
# ---------------------------------------------------------------------------- #
from datasets import Dataset, load_dataset, Sequence, ClassLabel
from tqdm import tqdm
from transformers import (AutoTokenizer)
import pandas as pd
import glob
import sys
import os
def get_root_dir():
path = os.path.dirname(__file__)
for _ in range(1):
path = os.path.dirname(path)
return path
sys.path.append(get_root_dir())
from src.utils.parse_data import parse_ast, parse_concept, parse_relation
# ---------------------------------------------------------------------------- #
# constants #
# ---------------------------------------------------------------------------- #
DEFAULT_DATA_DIR = os.path.join(get_root_dir(), "data")
NEW_LINE_CHAR = "\n"
MODEL_CHECKPOINT = "allenai/scibert_scivocab_uncased"
# ---------------------------------------------------------------------------- #
# Utils #
# ---------------------------------------------------------------------------- #
def masking(l): #to use if we only want concepts
if l >= 5:
if l % 2 == 0:
return 6
else :
return 5
else :
return l
def get_simple_tokenize(tokenizer):
def simple_tokenize(row):
tokens = tokenizer(row["text"], return_offsets_mapping=True)
return tokens
return simple_tokenize
def get_generate_row_labels(tokenizer, label_list, available_labels, verbose=False):
def generate_row_labels(row, verbose=verbose):
""" Given a row from the consolidated `Ade_corpus_v2_drug_ade_relation` dataset,
generates BIO tags for drug and effect entities.
"""
text = row["text"]
labels = []
label = "O"
prefix = ""
# while iterating through tokens, increment to traverse all drug and effect spans
label_index = {l : 0 for l in available_labels}
tokens = tokenizer(text, return_offsets_mapping=True)
for n in range(len(tokens["input_ids"])):
offset_start, offset_end = tokens["offset_mapping"][n]
# should only happen for [CLS] and [SEP]
if offset_end - offset_start == 0:
labels.append(-100)
continue
for l in available_labels :
if label_index[l] < len(row[l+"_indices_start"]) and offset_start == row[l+"_indices_start"][label_index[l]]:
label = l.upper()
prefix = "B-"
break
# labels.append(masking(label_list.index(f"{prefix}{label}")))
labels.append(label_list.index(f"{prefix}{label}"))
for l in available_labels :
if label_index[l] < len(row[l+"_indices_end"]) and offset_end == row[l+"_indices_end"][label_index[l]]:
label = "O"
prefix = ""
label_index[l] += 1
break
# need to transition "inside" if we just entered an entity
if prefix == "B-":
prefix = "I-"
if verbose:
print(row)
orig = tokenizer.convert_ids_to_tokens(tokens["input_ids"])
for n in range(len(labels)):
print(orig[n], labels[n])
tokens["labels"] = labels
return tokens
return generate_row_labels
# ---------------------------------------------------------------------------- #
# Main #
# ---------------------------------------------------------------------------- #
class Get_and_process_data:
def __init__(self, tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT),
data_path = DEFAULT_DATA_DIR, train_split = 0.8, add_unlabeled=True):
self.labels = ["test", "treatment",
"present", "absent", "possible", "conditional",
"hypothetical", "associated_with_someone_else"]
self.train_data_path = os.path.join(data_path, "train")
self.test_data_path = os.path.join(data_path, "test")
self.ast_folder_name = "ast"
self.concept_folder_name = "concept"
self.rel_folder_name = "rel"
self.txt_folder_name = "txt"
self.train_split = train_split
self.tokenizer = tokenizer
self.label_list = None
self.add_unlabeled = add_unlabeled
def load_raw_test(self):
text_files = glob.glob(self.test_data_path + os.sep + self.txt_folder_name + os.sep + "*.txt")
filename = ""
lines = pd.DataFrame()
for file in tqdm(text_files, "Loading raw text for test"):
with open(file, 'r') as f:
text = f.read()
filename = file.split("/")[-1].split(".")[0]
for i, line in enumerate(text.split(NEW_LINE_CHAR)):
lines = lines.append(pd.DataFrame({"text": [line], "filename": [filename], "line": i+1}), ignore_index=True)
preproc_data = {}
for i, row in tqdm(lines.iterrows(), "Formatting test data"):
line_id = row["filename"] + "_" + str(row["line"])
preproc_data[line_id] = {
"text": row["text"],
"row" : row["line"],
"filename": row["filename"],
}
# -------------------- Save and make hugging face dataset -------------------- #
preproc_df = pd.DataFrame(preproc_data).T
preproc_df.to_json(os.path.join(self.test_data_path, "dataset.jsonl"), orient="records", lines=True)
dataset = load_dataset("json", data_files={"test" : os.path.join(self.test_data_path, "dataset.jsonl")})
# ------------------ Creating the right format and tokenize ------------------ #
label_list = ['O']
for label in self.labels:
label_list.append("B-"+label.upper())
label_list.append("I-"+label.upper())
custom_seq = Sequence(feature=ClassLabel(num_classes=len(label_list),
names=label_list,
names_file=None, id=None), length=-1, id=None)
dataset["test"].features["ner_tags"] = custom_seq
labeled_dataset = dataset.map(get_simple_tokenize(self.tokenizer))
return labeled_dataset
def load_parse(self):
"""
Output :
ast_concept_df :
concept_text start_line start_word_number end_line end_word_number ast_con_label filename
0 pain 55 10 55 10 hypothetical 018636330_DH.txt
1 hyperlipidemia 29 4 29 4 present 018636330_DH
raw_text :
text
filename
018636330_DH 018636330 DH\n5425710\n123524\n0144918\n6/2/20...
"""
text_files = glob.glob(self.train_data_path + os.sep + self.txt_folder_name + os.sep + "*.txt")
filename = ""
raw_files = pd.DataFrame()
unlabeled_lines = pd.DataFrame()
for file in tqdm(text_files, "Loading raw text"):
with open(file, 'r') as f:
text = f.read()
filename = file.split("/")[-1].split(".")[0]
ast = parse_ast(self.train_data_path + os.sep + self.ast_folder_name + os.sep + filename + ".ast")
concept = parse_concept(self.train_data_path + os.sep + self.concept_folder_name + os.sep + filename + ".con")
rel = parse_relation(self.train_data_path + os.sep + self.rel_folder_name + os.sep + filename + ".rel")
raw_files = raw_files.append(pd.DataFrame({"text": [text], "filename": [filename] , "concept": [concept], "ast": [ast], "rel": [rel]}), ignore_index=True)
# -------------------- known lines are the ones in Concept ------------------- #
if self.add_unlabeled:
known_lines = concept["start_line"]
for i, line in enumerate(text.split(NEW_LINE_CHAR)):
if not i+1 in known_lines:
unlabeled_lines = unlabeled_lines.append(pd.DataFrame({"text": [line], "filename": [filename], "line": i+1}), ignore_index=True)
raw_text = raw_files[["text", "filename"]].set_index("filename")
# ---------------------------- Mixing the 2 tasks ---------------------------- #
ast_concept_df = pd.DataFrame()
for i, file in tqdm(raw_files.iterrows(), "Processing raw text"):
ast_dict = file["ast"]
concept_dict = file["concept"]
tmp_ast = pd.DataFrame(ast_dict)
tmp_ast = tmp_ast.drop(columns=["concept_type"])
tmp_ast = tmp_ast.rename(columns={"assertion_type": "ast_con_label"})
#Only concepts with not "problem"
tmp_concept = pd.DataFrame(concept_dict)
if len(tmp_ast) > 0:
assert(
tmp_concept[tmp_concept["concept_type"]=="problem"]["concept_text"].reset_index(drop=True).equals(
tmp_ast["concept_text"]
)
), "Concepts with problem type are not the same as assertions"
tmp_concept = tmp_concept.rename(columns={"concept_type": "ast_con_label"})
tmp = tmp_concept[tmp_concept["ast_con_label"]=="problem"].reset_index(drop=True)
tmp["ast_con_label"] = tmp_ast["ast_con_label"]
tmp_ast = tmp
tmp_concept = tmp_concept[tmp_concept["ast_con_label"] != "problem"]
tmp_ast["filename"] = file["filename"]
tmp_concept["filename"] = file["filename"]
if len(tmp_ast) > 0:
ast_concept_df = ast_concept_df.append(tmp_ast, ignore_index=True)
if len(tmp_concept) > 0:
ast_concept_df = ast_concept_df.append(tmp_concept, ignore_index=True)
#cols = concept_text, start_line, start_word_number, end_line, end_word_number, ast_con_label, filename
return ast_concept_df, raw_text, unlabeled_lines
def format(self, ast_concept_df : pd.DataFrame, raw_text : pd.DataFrame, unlabeled_lines : pd.DataFrame):
preproc_data = {}
# -------------------------- Add the unlabeled_lines ------------------------- #
for i, row in tqdm(unlabeled_lines.iterrows(), "Adding unlabeled lines"):
line_id = row["filename"] + "_" + str(row["line"])
preproc_data[line_id] = {
"text": row["text"],
"row" : row["line"],
"filename": row["filename"],
}
for label in self.labels:
preproc_data[line_id][label] = []
# use sets because the indices can repeat for various reasons
preproc_data[line_id][label+"_indices_start"] = set()
preproc_data[line_id][label+"_indices_end"] = set()
for i, row in tqdm(ast_concept_df.iterrows(), "Formatting dataset"):
filename = row["filename"]
text = raw_text.loc[filename]["text"]
# text Normalization
text = text.lower()
line = text.split(NEW_LINE_CHAR)[row["start_line"] - 1] # NOTE: we assume that start_line == end_line
line = " ".join(line.split()) # remove multiple spaces
row["concept_text"] = " ".join(row["concept_text"].split()) # remove multiple spaces
# find character index start and end of concept
start_char_index = len(" ".join(line.split()[: row["start_word_number"]])) # number of chars before concept
if start_char_index > 0:
start_char_index += 1
end_char_index = start_char_index + len(row["concept_text"])
assert (
line[start_char_index:end_char_index] == row["concept_text"]
), f"concept_text doesn't match the found indexes. '{line[start_char_index:end_char_index]}' != '{row['concept_text']}'"
line_id = filename + "_" + str(row["start_line"])
if line_id not in preproc_data:
preproc_data[line_id] = {
"text": line,
"row" : row["start_line"],
"filename": row["filename"],
}
for label in self.labels:
preproc_data[line_id][label] = []
# use sets because the indices can repeat for various reasons
preproc_data[line_id][label+"_indices_start"] = set()
preproc_data[line_id][label+"_indices_end"] = set()
for label in self.labels:
if row["ast_con_label"] == label:
preproc_data[line_id][label].append(row["concept_text"])
preproc_data[line_id][label+"_indices_start"].add(start_char_index)
preproc_data[line_id][label+"_indices_end"].add(end_char_index)
break
preproc_df = pd.DataFrame(preproc_data).T
for label in self.labels:
preproc_df[label+"_indices_start"] = preproc_df[label+"_indices_start"].apply(list).apply(sorted)
preproc_df[label+"_indices_end"] = preproc_df[label+"_indices_end"].apply(list).apply(sorted)
preproc_df.to_json(os.path.join(self.train_data_path, "dataset.jsonl"), orient="records", lines=True)
dataset = load_dataset("json", data_files=os.path.join(self.train_data_path, "dataset.jsonl"))
dataset = dataset["train"].train_test_split(train_size=self.train_split)
# ---------------------------------- rename ---------------------------------- #
dataset["val"] = dataset["test"]
return dataset
def token_labeling(self, dataset : Dataset):
label_list = ['O']
for label in self.labels:
label_list.append("B-"+label.upper())
label_list.append("I-"+label.upper())
custom_seq = Sequence(feature=ClassLabel(num_classes=len(label_list),
names=label_list,
names_file=None, id=None), length=-1, id=None)
dataset["train"].features["ner_tags"] = custom_seq
dataset["val"].features["ner_tags"] = custom_seq
labeled_dataset = dataset.map(get_generate_row_labels(self.tokenizer, label_list, self.labels))
self.label_list = label_list
# ----------------------------- Adding test data ----------------------------- #
test_dataset = self.load_raw_test()
labeled_dataset["test"] = test_dataset["test"]
return labeled_dataset
def get_label_list(self):
return self.label_list
def get_dataset(self):
dataset = self.format(*self.load_parse())
return self.token_labeling(dataset)
def get_default_dataset():
return Get_and_process_data().get_dataset()
if __name__ == "__main__":
print(get_default_dataset())
