def __init__(self, squad_model_config: str, vocab_file: str, do_lower_case: bool, max_seq_length: int = 512, batch_size: int = 10, lang: str = 'en', **kwargs) -> None: config = json.load(open(squad_model_config)) config['chainer']['pipe'][0]['max_seq_length'] = max_seq_length self.model = build_model(config) self.max_seq_length = max_seq_length if Path(vocab_file).is_file(): vocab_file = str(expand_path(vocab_file)) self.tokenizer = AutoTokenizer(vocab_file=vocab_file, do_lower_case=do_lower_case) else: self.tokenizer = AutoTokenizer.from_pretrained( vocab_file, do_lower_case=do_lower_case) self.batch_size = batch_size if lang == 'en': from nltk import sent_tokenize self.sent_tokenizer = sent_tokenize elif lang == 'ru': from ru_sent_tokenize import ru_sent_tokenize self.sent_tokenizer = ru_sent_tokenize else: raise RuntimeError('en and ru languages are supported only')
def predict(self, sentence: str) -> torch.Tensor: tokenized_input = AutoTokenizer(sentence) output = self.model(tokenized_input["input_ids"], tokenized_input["attention_mask"]) logit = output.logits pred_idx = torch.argmax(logit, dim=-1) return self.idx_2_label[pred_idx]
def main(test_dir_pattern): test_files = glob.glob(test_dir_pattern) tokenizer = AutoTokenizer(config.CHECKPOINT) ner_model = NERModel() ner_model.load(config.MODEL_CHECKPOINT) bert = ner_model.bert.eval().cuda() last_layer = ner_model.last_layer.eval().cuda() mapping = {} submission_mapping = defaultdict(list) doc_index = 0 n_files = len(test_files) for i in range(0, n_files, 50): files_to_processed = test_files[i:i + 50] test_examples = [] for file in files_to_processed: doc = Document.read_from_path(file) sentences = doc.get_all_sentences() mapping[doc_index] = doc.document_id for sentence in sentences: te = TestExample(doc_index, sentence) test_examples.append(te) doc_index += 1 dataset = NERTestDataset(test_examples, tokenizer) data_loader = DataLoader(dataset, batch_size=config.TEST_BATCH_SIZE) for batch in data_loader: o = bert( input_ids=batch['input_ids'].to('cuda'), attention_mask=batch['attention_mask'].to('cuda'), token_type_ids=batch['token_type_ids'].to('cuda'), ) o = last_layer(o['last_hidden_state']) o = o.detach() o = o.argmax(axis=-1) probabilities = o.max(axis=-1) for j, sentence_prediction in enumerate(o): last_valid = batch['attention_mask'][j].argmin().item() candidates = LabelEncoding.extract_candidates( sentence_prediction, probabilities[j], last_valid, ) if len(candidates) == 0: continue id_ = batch['id'][j].item() ids = batch['input_ids'][j] for start_index, end_index in candidates: submission_mapping[mapping[id_]].append( tokenizer.decode(ids[start_index:end_index])) return submission_mapping
def __init__(self, vocab_file: str, do_lower_case: bool = True, max_seq_length: int = 512, return_tokens: bool = False, **kwargs) -> None: self.max_seq_length = max_seq_length self.return_tokens = return_tokens if Path(vocab_file).is_file(): vocab_file = str(expand_path(vocab_file)) self.tokenizer = AutoTokenizer(vocab_file=vocab_file, do_lower_case=do_lower_case) else: self.tokenizer = AutoTokenizer.from_pretrained(vocab_file, do_lower_case=do_lower_case)
def __init__(self, vocab_file: str, do_lower_case: bool = False, max_seq_length: int = 512, max_subword_length: int = None, token_masking_prob: float = 0.0, provide_subword_tags: bool = False, subword_mask_mode: str = "first", **kwargs): self._re_tokenizer = re.compile(r"[\w']+|[^\w ]") self.provide_subword_tags = provide_subword_tags self.mode = kwargs.get('mode') self.max_seq_length = max_seq_length self.max_subword_length = max_subword_length self.subword_mask_mode = subword_mask_mode if Path(vocab_file).is_file(): vocab_file = str(expand_path(vocab_file)) self.tokenizer = AutoTokenizer(vocab_file=vocab_file, do_lower_case=do_lower_case) else: self.tokenizer = AutoTokenizer.from_pretrained(vocab_file, do_lower_case=do_lower_case) self.token_masking_prob = token_masking_prob