shengz's picture
Add the example usage.
e3ef0b9
raw
history blame
No virus
9.69 kB
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from typing import List, Dict
import os
import time
import logging
import json
import gzip
from dataclasses import dataclass, field
import torch
from torch import Tensor as T
from transformers import PreTrainedTokenizer
logger = logging.getLogger()
@dataclass
class Mention:
cui: str
start: int
end: int
text: str
types: str
@dataclass
class ContextualMention:
mention: str
cuis: List[str]
ctx_l: str
ctx_r: str
def to_tensor(self, tokenizer: PreTrainedTokenizer, max_length: int) -> T:
ctx_l_ids = tokenizer.encode(
text=self.ctx_l,
add_special_tokens=False,
max_length=max_length,
truncation=True,
)
ctx_r_ids = tokenizer.encode(
text=self.ctx_r,
add_special_tokens=False,
max_length=max_length,
truncation=True,
)
mention_ids = tokenizer.encode(
text=self.mention,
add_special_tokens=False,
max_length=max_length,
truncation=True,
)
# Concatenate context and mention to the max length.
token_ids = tokenizer.convert_tokens_to_ids(['<ENT>']) + mention_ids \
+ tokenizer.convert_tokens_to_ids(['</ENT>'])
max_ctx_len = max_length - len(token_ids) - 2 # Exclude [CLS] and [SEP]
max_ctx_l_len = max_ctx_len // 2
max_ctx_r_len = max_ctx_len - max_ctx_l_len
if len(ctx_l_ids) < max_ctx_l_len and len(ctx_r_ids) < max_ctx_r_len:
token_ids = ctx_l_ids + token_ids + ctx_r_ids
elif len(ctx_l_ids) >= max_ctx_l_len and len(ctx_r_ids) >= max_ctx_r_len:
token_ids = ctx_l_ids[-max_ctx_l_len:] + token_ids \
+ ctx_r_ids[:max_ctx_r_len]
elif len(ctx_l_ids) >= max_ctx_l_len:
ctx_l_len = max_ctx_len - len(ctx_r_ids)
token_ids = ctx_l_ids[-ctx_l_len:] + token_ids + ctx_r_ids
else:
ctx_r_len = max_ctx_len - len(ctx_l_ids)
token_ids = ctx_l_ids + token_ids + ctx_r_ids[:ctx_r_len]
token_ids = [tokenizer.cls_token_id] + token_ids
# The above snippet doesn't guarantee the max length limit.
token_ids = token_ids[:max_length - 1] + [tokenizer.sep_token_id]
if len(token_ids) < max_length:
token_ids = token_ids + [tokenizer.pad_token_id] * (max_length - len(token_ids))
return torch.tensor(token_ids)
@dataclass
class Document:
id: str = None
title: str = None
abstract: str = None
mentions: List[Mention] = field(default_factory=list)
def concatenate_text(self) -> str:
return ' '.join([self.title, self.abstract])
@classmethod
def from_PubTator(cls, path: str, split_path_prefix: str) -> Dict[str, List]:
docs = []
with gzip.open(path, 'rb') as f:
for b in f.read().decode().strip().split('\n\n'):
d = cls()
s = ''
for i, ln in enumerate(b.split('\n')):
if i == 0:
id, type, text = ln.strip().split('|', 2)
assert type == 't'
d.id, d.title = id, text
elif i == 1:
id, type, text = ln.strip().split('|', 2)
assert type == 'a'
assert d.id == id
d.abstract = text
s = d.concatenate_text()
else:
items = ln.strip().split('\t')
assert d.id == items[0]
cui = items[5].split('UMLS:')[-1]
assert len(cui) == 8, breakpoint()
m = Mention(
cui=cui,
start=int(items[1]),
end=int(items[2]),
text=items[3],
types=items[4].split(',')
)
assert m.text == s[m.start: m.end]
d.mentions.append(m)
docs.append(d)
dataset = split_dataset(docs, split_path_prefix)
print_dataset_stats(dataset)
return dataset
def to_contextual_mentions(self, max_length: int = 64) -> List[ContextualMention]:
text = self.concatenate_text()
mentions = []
for m in self.mentions:
assert m.text == text[m.start:m.end]
# Context
ctx_l, ctx_r = text[:m.start].strip().split(), text[m.end:].strip().split()
ctx_l, ctx_r = ' '.join(ctx_l[-max_length:]), ' '.join(ctx_r[:max_length])
cm = ContextualMention(
mention=m.text,
cuis=[m.cui],
ctx_l=ctx_l,
ctx_r=ctx_r,
)
mentions.append(cm)
return mentions
def split_dataset(docs: List, split_path_prefix: str) -> Dict[str, List]:
split_kv = {'train': 'trng', 'dev': 'dev', 'test': 'test'}
id_to_split = {}
dataset = {}
for k, v in split_kv.items():
dataset[k] = []
path = split_path_prefix + v + '.txt'
for i in open(path, encoding='utf-8').read().strip().split('\n'):
assert i not in id_to_split, breakpoint()
id_to_split[i] = k
for doc in docs:
split = id_to_split[doc.id]
dataset[split].append(doc)
return dataset
def print_dataset_stats(dataset: Dict[str, List[Document]]) -> None:
all_docs = []
for v in dataset.values():
all_docs.extend(v)
for split, docs in {'all': all_docs, **dataset}.items():
logger.info(f"***** {split} *****")
logger.info(f"Documents: {len(docs)}")
logger.info(f"Mentions: {sum(len(d.mentions) for d in docs)}")
cuis = set()
for d in docs:
for m in d.mentions:
cuis.add(m.cui)
logger.info(f"Mentioned concepts: {len(cuis)}")
class MedMentionsDataset(torch.utils.data.Dataset):
def __init__(self, dataset_path: str, split: str) -> None:
super().__init__()
self.dataset_path = dataset_path
self.docs = Document.from_PubTator(
path=os.path.join(self.dataset_path, 'corpus_pubtator.txt.gz'),
split_path_prefix=os.path.join(self.dataset_path, 'corpus_pubtator_pmids_')
)[split]
self.mentions = []
self.name_to_cuis = {}
self._post_init()
def _post_init(self):
for d in self.docs:
self.mentions.extend(d.to_contextual_mentions())
for m in self.mentions:
if m.mention not in self.name_to_cuis:
self.name_to_cuis[m.mention] = set()
self.name_to_cuis[m.mention].update(m.cuis)
def __getitem__(self, index: int) -> ContextualMention:
return self.mentions[index]
def __len__(self) -> int:
return len(self.mentions)
class PreprocessedDataset(torch.utils.data.Dataset):
def __init__(self, dataset_path: str) -> None:
super().__init__()
self.file = dataset_path
self.data = []
self.load_data()
def load_data(self) -> None:
with open(self.file, encoding='utf-8') as f:
logger.info("Reading file %s" % self.file)
for ln in f:
if ln.strip():
self.data.append(json.loads(ln))
logger.info("Loaded data size: {}".format(len(self.data)))
def __getitem__(self, index: int) -> ContextualMention:
d = self.data[index]
return ContextualMention(
ctx_l=d['context_left'],
ctx_r=d['context_right'],
mention=d['mention'],
cuis=d['cuis'],
)
def __len__(self) -> int:
return len(self.data)
def generate_vectors(
encoder: torch.nn.Module,
tokenizer: PreTrainedTokenizer,
dataset: torch.utils.data.Dataset,
batch_size: int,
max_length: int,
is_prototype: bool = False,
):
n = len(dataset)
total = 0
results = []
start_time = time.time()
logger.info("Start encoding...")
for i, batch_start in enumerate(range(0, n, batch_size)):
batch = [dataset[i] for i in range(batch_start, min(n, batch_start + batch_size))]
batch_token_tensors = [m.to_tensor(tokenizer, max_length) for m in batch]
ids_batch = torch.stack(batch_token_tensors, dim=0).cuda()
seg_batch = torch.zeros_like(ids_batch)
attn_mask = (ids_batch != tokenizer.pad_token_id)
with torch.inference_mode():
out = encoder(
input_ids=ids_batch,
token_type_ids=seg_batch,
attention_mask=attn_mask
)
out = out[0][:, 0, :]
out = out.cpu()
num_mentions = out.size(0)
total += num_mentions
if is_prototype:
meta_batch = [{'cuis': m.cuis} for m in batch]
assert len(meta_batch) == num_mentions
results.extend([(meta_batch[i], out[i].view(-1).numpy()) for i in range(num_mentions)])
else:
results.extend(out.cpu().split(1, dim=0))
if (i + 1) % 10 == 0:
eta = (n - total) * (time.time() - start_time) / 60 / total
logger.info(f"Batch={i + 1}, Encoded mentions={total}, ETA={eta:.1f}m")
assert len(results) == n
logger.info(f"Total encoded mentions={n}")
if not is_prototype:
results = torch.cat(results, dim=0)
return results