squad.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845
  1. # Copyright 2020 The HuggingFace Team. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import json
  15. import os
  16. from functools import partial
  17. from multiprocessing import Pool, cpu_count
  18. from multiprocessing.pool import ThreadPool
  19. from typing import Optional
  20. import numpy as np
  21. from tqdm import tqdm
  22. from ...models.bert.tokenization_bert import whitespace_tokenize
  23. from ...tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase, TruncationStrategy
  24. from ...utils import is_tf_available, is_torch_available, is_torch_hpu_available, logging
  25. from .utils import DataProcessor
  26. # Store the tokenizers which insert 2 separators tokens
  27. MULTI_SEP_TOKENS_TOKENIZERS_SET = {"roberta", "camembert", "bart", "mpnet"}
  28. if is_torch_available():
  29. import torch
  30. from torch.utils.data import TensorDataset
  31. if is_tf_available():
  32. import tensorflow as tf
  33. logger = logging.get_logger(__name__)
  34. def _improve_answer_span(doc_tokens, input_start, input_end, tokenizer, orig_answer_text):
  35. """Returns tokenized answer spans that better match the annotated answer."""
  36. tok_answer_text = " ".join(tokenizer.tokenize(orig_answer_text))
  37. for new_start in range(input_start, input_end + 1):
  38. for new_end in range(input_end, new_start - 1, -1):
  39. text_span = " ".join(doc_tokens[new_start : (new_end + 1)])
  40. if text_span == tok_answer_text:
  41. return (new_start, new_end)
  42. return (input_start, input_end)
  43. def _check_is_max_context(doc_spans, cur_span_index, position):
  44. """Check if this is the 'max context' doc span for the token."""
  45. best_score = None
  46. best_span_index = None
  47. for span_index, doc_span in enumerate(doc_spans):
  48. end = doc_span.start + doc_span.length - 1
  49. if position < doc_span.start:
  50. continue
  51. if position > end:
  52. continue
  53. num_left_context = position - doc_span.start
  54. num_right_context = end - position
  55. score = min(num_left_context, num_right_context) + 0.01 * doc_span.length
  56. if best_score is None or score > best_score:
  57. best_score = score
  58. best_span_index = span_index
  59. return cur_span_index == best_span_index
  60. def _new_check_is_max_context(doc_spans, cur_span_index, position):
  61. """Check if this is the 'max context' doc span for the token."""
  62. # if len(doc_spans) == 1:
  63. # return True
  64. best_score = None
  65. best_span_index = None
  66. for span_index, doc_span in enumerate(doc_spans):
  67. end = doc_span["start"] + doc_span["length"] - 1
  68. if position < doc_span["start"]:
  69. continue
  70. if position > end:
  71. continue
  72. num_left_context = position - doc_span["start"]
  73. num_right_context = end - position
  74. score = min(num_left_context, num_right_context) + 0.01 * doc_span["length"]
  75. if best_score is None or score > best_score:
  76. best_score = score
  77. best_span_index = span_index
  78. return cur_span_index == best_span_index
  79. def _is_whitespace(c):
  80. if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
  81. return True
  82. return False
  83. def squad_convert_example_to_features(
  84. example, max_seq_length, doc_stride, max_query_length, padding_strategy, is_training
  85. ):
  86. features = []
  87. if is_training and not example.is_impossible:
  88. # Get start and end position
  89. start_position = example.start_position
  90. end_position = example.end_position
  91. # If the answer cannot be found in the text, then skip this example.
  92. actual_text = " ".join(example.doc_tokens[start_position : (end_position + 1)])
  93. cleaned_answer_text = " ".join(whitespace_tokenize(example.answer_text))
  94. if actual_text.find(cleaned_answer_text) == -1:
  95. logger.warning(f"Could not find answer: '{actual_text}' vs. '{cleaned_answer_text}'")
  96. return []
  97. tok_to_orig_index = []
  98. orig_to_tok_index = []
  99. all_doc_tokens = []
  100. for i, token in enumerate(example.doc_tokens):
  101. orig_to_tok_index.append(len(all_doc_tokens))
  102. if tokenizer.__class__.__name__ in [
  103. "RobertaTokenizer",
  104. "LongformerTokenizer",
  105. "BartTokenizer",
  106. "RobertaTokenizerFast",
  107. "LongformerTokenizerFast",
  108. "BartTokenizerFast",
  109. ]:
  110. sub_tokens = tokenizer.tokenize(token, add_prefix_space=True)
  111. else:
  112. sub_tokens = tokenizer.tokenize(token)
  113. for sub_token in sub_tokens:
  114. tok_to_orig_index.append(i)
  115. all_doc_tokens.append(sub_token)
  116. if is_training and not example.is_impossible:
  117. tok_start_position = orig_to_tok_index[example.start_position]
  118. if example.end_position < len(example.doc_tokens) - 1:
  119. tok_end_position = orig_to_tok_index[example.end_position + 1] - 1
  120. else:
  121. tok_end_position = len(all_doc_tokens) - 1
  122. (tok_start_position, tok_end_position) = _improve_answer_span(
  123. all_doc_tokens, tok_start_position, tok_end_position, tokenizer, example.answer_text
  124. )
  125. spans = []
  126. truncated_query = tokenizer.encode(
  127. example.question_text, add_special_tokens=False, truncation=True, max_length=max_query_length
  128. )
  129. # Tokenizers who insert 2 SEP tokens in-between <context> & <question> need to have special handling
  130. # in the way they compute mask of added tokens.
  131. tokenizer_type = type(tokenizer).__name__.replace("Tokenizer", "").lower()
  132. sequence_added_tokens = (
  133. tokenizer.model_max_length - tokenizer.max_len_single_sentence + 1
  134. if tokenizer_type in MULTI_SEP_TOKENS_TOKENIZERS_SET
  135. else tokenizer.model_max_length - tokenizer.max_len_single_sentence
  136. )
  137. sequence_pair_added_tokens = tokenizer.model_max_length - tokenizer.max_len_sentences_pair
  138. span_doc_tokens = all_doc_tokens
  139. while len(spans) * doc_stride < len(all_doc_tokens):
  140. # Define the side we want to truncate / pad and the text/pair sorting
  141. if tokenizer.padding_side == "right":
  142. texts = truncated_query
  143. pairs = span_doc_tokens
  144. truncation = TruncationStrategy.ONLY_SECOND.value
  145. else:
  146. texts = span_doc_tokens
  147. pairs = truncated_query
  148. truncation = TruncationStrategy.ONLY_FIRST.value
  149. encoded_dict = tokenizer.encode_plus( # TODO(thom) update this logic
  150. texts,
  151. pairs,
  152. truncation=truncation,
  153. padding=padding_strategy,
  154. max_length=max_seq_length,
  155. return_overflowing_tokens=True,
  156. stride=max_seq_length - doc_stride - len(truncated_query) - sequence_pair_added_tokens,
  157. return_token_type_ids=True,
  158. )
  159. paragraph_len = min(
  160. len(all_doc_tokens) - len(spans) * doc_stride,
  161. max_seq_length - len(truncated_query) - sequence_pair_added_tokens,
  162. )
  163. if tokenizer.pad_token_id in encoded_dict["input_ids"]:
  164. if tokenizer.padding_side == "right":
  165. non_padded_ids = encoded_dict["input_ids"][: encoded_dict["input_ids"].index(tokenizer.pad_token_id)]
  166. else:
  167. last_padding_id_position = (
  168. len(encoded_dict["input_ids"]) - 1 - encoded_dict["input_ids"][::-1].index(tokenizer.pad_token_id)
  169. )
  170. non_padded_ids = encoded_dict["input_ids"][last_padding_id_position + 1 :]
  171. else:
  172. non_padded_ids = encoded_dict["input_ids"]
  173. tokens = tokenizer.convert_ids_to_tokens(non_padded_ids)
  174. token_to_orig_map = {}
  175. for i in range(paragraph_len):
  176. index = len(truncated_query) + sequence_added_tokens + i if tokenizer.padding_side == "right" else i
  177. token_to_orig_map[index] = tok_to_orig_index[len(spans) * doc_stride + i]
  178. encoded_dict["paragraph_len"] = paragraph_len
  179. encoded_dict["tokens"] = tokens
  180. encoded_dict["token_to_orig_map"] = token_to_orig_map
  181. encoded_dict["truncated_query_with_special_tokens_length"] = len(truncated_query) + sequence_added_tokens
  182. encoded_dict["token_is_max_context"] = {}
  183. encoded_dict["start"] = len(spans) * doc_stride
  184. encoded_dict["length"] = paragraph_len
  185. spans.append(encoded_dict)
  186. if "overflowing_tokens" not in encoded_dict or (
  187. "overflowing_tokens" in encoded_dict and len(encoded_dict["overflowing_tokens"]) == 0
  188. ):
  189. break
  190. span_doc_tokens = encoded_dict["overflowing_tokens"]
  191. for doc_span_index in range(len(spans)):
  192. for j in range(spans[doc_span_index]["paragraph_len"]):
  193. is_max_context = _new_check_is_max_context(spans, doc_span_index, doc_span_index * doc_stride + j)
  194. index = (
  195. j
  196. if tokenizer.padding_side == "left"
  197. else spans[doc_span_index]["truncated_query_with_special_tokens_length"] + j
  198. )
  199. spans[doc_span_index]["token_is_max_context"][index] = is_max_context
  200. for span in spans:
  201. # Identify the position of the CLS token
  202. cls_index = span["input_ids"].index(tokenizer.cls_token_id)
  203. # p_mask: mask with 1 for token than cannot be in the answer (0 for token which can be in an answer)
  204. # Original TF implementation also keep the classification token (set to 0)
  205. p_mask = np.ones_like(span["token_type_ids"])
  206. if tokenizer.padding_side == "right":
  207. p_mask[len(truncated_query) + sequence_added_tokens :] = 0
  208. else:
  209. p_mask[-len(span["tokens"]) : -(len(truncated_query) + sequence_added_tokens)] = 0
  210. pad_token_indices = np.where(np.atleast_1d(span["input_ids"] == tokenizer.pad_token_id))
  211. special_token_indices = np.asarray(
  212. tokenizer.get_special_tokens_mask(span["input_ids"], already_has_special_tokens=True)
  213. ).nonzero()
  214. p_mask[pad_token_indices] = 1
  215. p_mask[special_token_indices] = 1
  216. # Set the cls index to 0: the CLS index can be used for impossible answers
  217. p_mask[cls_index] = 0
  218. span_is_impossible = example.is_impossible
  219. start_position = 0
  220. end_position = 0
  221. if is_training and not span_is_impossible:
  222. # For training, if our document chunk does not contain an annotation
  223. # we throw it out, since there is nothing to predict.
  224. doc_start = span["start"]
  225. doc_end = span["start"] + span["length"] - 1
  226. out_of_span = False
  227. if not (tok_start_position >= doc_start and tok_end_position <= doc_end):
  228. out_of_span = True
  229. if out_of_span:
  230. start_position = cls_index
  231. end_position = cls_index
  232. span_is_impossible = True
  233. else:
  234. if tokenizer.padding_side == "left":
  235. doc_offset = 0
  236. else:
  237. doc_offset = len(truncated_query) + sequence_added_tokens
  238. start_position = tok_start_position - doc_start + doc_offset
  239. end_position = tok_end_position - doc_start + doc_offset
  240. features.append(
  241. SquadFeatures(
  242. span["input_ids"],
  243. span["attention_mask"],
  244. span["token_type_ids"],
  245. cls_index,
  246. p_mask.tolist(),
  247. example_index=0, # Can not set unique_id and example_index here. They will be set after multiple processing.
  248. unique_id=0,
  249. paragraph_len=span["paragraph_len"],
  250. token_is_max_context=span["token_is_max_context"],
  251. tokens=span["tokens"],
  252. token_to_orig_map=span["token_to_orig_map"],
  253. start_position=start_position,
  254. end_position=end_position,
  255. is_impossible=span_is_impossible,
  256. qas_id=example.qas_id,
  257. )
  258. )
  259. return features
  260. def squad_convert_example_to_features_init(tokenizer_for_convert: PreTrainedTokenizerBase):
  261. global tokenizer
  262. tokenizer = tokenizer_for_convert
  263. def squad_convert_examples_to_features(
  264. examples,
  265. tokenizer,
  266. max_seq_length,
  267. doc_stride,
  268. max_query_length,
  269. is_training,
  270. padding_strategy="max_length",
  271. return_dataset=False,
  272. threads=1,
  273. tqdm_enabled=True,
  274. ):
  275. """
  276. Converts a list of examples into a list of features that can be directly given as input to a model. It is
  277. model-dependant and takes advantage of many of the tokenizer's features to create the model's inputs.
  278. Args:
  279. examples: list of [`~data.processors.squad.SquadExample`]
  280. tokenizer: an instance of a child of [`PreTrainedTokenizer`]
  281. max_seq_length: The maximum sequence length of the inputs.
  282. doc_stride: The stride used when the context is too large and is split across several features.
  283. max_query_length: The maximum length of the query.
  284. is_training: whether to create features for model evaluation or model training.
  285. padding_strategy: Default to "max_length". Which padding strategy to use
  286. return_dataset: Default False. Either 'pt' or 'tf'.
  287. if 'pt': returns a torch.data.TensorDataset, if 'tf': returns a tf.data.Dataset
  288. threads: multiple processing threads.
  289. Returns:
  290. list of [`~data.processors.squad.SquadFeatures`]
  291. Example:
  292. ```python
  293. processor = SquadV2Processor()
  294. examples = processor.get_dev_examples(data_dir)
  295. features = squad_convert_examples_to_features(
  296. examples=examples,
  297. tokenizer=tokenizer,
  298. max_seq_length=args.max_seq_length,
  299. doc_stride=args.doc_stride,
  300. max_query_length=args.max_query_length,
  301. is_training=not evaluate,
  302. )
  303. ```"""
  304. threads = min(threads, cpu_count())
  305. pool_cls = ThreadPool if is_torch_hpu_available() else Pool
  306. with pool_cls(threads, initializer=squad_convert_example_to_features_init, initargs=(tokenizer,)) as p:
  307. annotate_ = partial(
  308. squad_convert_example_to_features,
  309. max_seq_length=max_seq_length,
  310. doc_stride=doc_stride,
  311. max_query_length=max_query_length,
  312. padding_strategy=padding_strategy,
  313. is_training=is_training,
  314. )
  315. features = list(
  316. tqdm(
  317. p.imap(annotate_, examples, chunksize=32),
  318. total=len(examples),
  319. desc="convert squad examples to features",
  320. disable=not tqdm_enabled,
  321. )
  322. )
  323. new_features = []
  324. unique_id = 1000000000
  325. example_index = 0
  326. for example_features in tqdm(
  327. features, total=len(features), desc="add example index and unique id", disable=not tqdm_enabled
  328. ):
  329. if not example_features:
  330. continue
  331. for example_feature in example_features:
  332. example_feature.example_index = example_index
  333. example_feature.unique_id = unique_id
  334. new_features.append(example_feature)
  335. unique_id += 1
  336. example_index += 1
  337. features = new_features
  338. del new_features
  339. if return_dataset == "pt":
  340. if not is_torch_available():
  341. raise RuntimeError("PyTorch must be installed to return a PyTorch dataset.")
  342. # Convert to Tensors and build dataset
  343. all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
  344. all_attention_masks = torch.tensor([f.attention_mask for f in features], dtype=torch.long)
  345. all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)
  346. all_cls_index = torch.tensor([f.cls_index for f in features], dtype=torch.long)
  347. all_p_mask = torch.tensor([f.p_mask for f in features], dtype=torch.float)
  348. all_is_impossible = torch.tensor([f.is_impossible for f in features], dtype=torch.float)
  349. if not is_training:
  350. all_feature_index = torch.arange(all_input_ids.size(0), dtype=torch.long)
  351. dataset = TensorDataset(
  352. all_input_ids, all_attention_masks, all_token_type_ids, all_feature_index, all_cls_index, all_p_mask
  353. )
  354. else:
  355. all_start_positions = torch.tensor([f.start_position for f in features], dtype=torch.long)
  356. all_end_positions = torch.tensor([f.end_position for f in features], dtype=torch.long)
  357. dataset = TensorDataset(
  358. all_input_ids,
  359. all_attention_masks,
  360. all_token_type_ids,
  361. all_start_positions,
  362. all_end_positions,
  363. all_cls_index,
  364. all_p_mask,
  365. all_is_impossible,
  366. )
  367. return features, dataset
  368. elif return_dataset == "tf":
  369. if not is_tf_available():
  370. raise RuntimeError("TensorFlow must be installed to return a TensorFlow dataset.")
  371. def gen():
  372. for i, ex in enumerate(features):
  373. if ex.token_type_ids is None:
  374. yield (
  375. {
  376. "input_ids": ex.input_ids,
  377. "attention_mask": ex.attention_mask,
  378. "feature_index": i,
  379. "qas_id": ex.qas_id,
  380. },
  381. {
  382. "start_positions": ex.start_position,
  383. "end_positions": ex.end_position,
  384. "cls_index": ex.cls_index,
  385. "p_mask": ex.p_mask,
  386. "is_impossible": ex.is_impossible,
  387. },
  388. )
  389. else:
  390. yield (
  391. {
  392. "input_ids": ex.input_ids,
  393. "attention_mask": ex.attention_mask,
  394. "token_type_ids": ex.token_type_ids,
  395. "feature_index": i,
  396. "qas_id": ex.qas_id,
  397. },
  398. {
  399. "start_positions": ex.start_position,
  400. "end_positions": ex.end_position,
  401. "cls_index": ex.cls_index,
  402. "p_mask": ex.p_mask,
  403. "is_impossible": ex.is_impossible,
  404. },
  405. )
  406. # Why have we split the batch into a tuple? PyTorch just has a list of tensors.
  407. if "token_type_ids" in tokenizer.model_input_names:
  408. train_types = (
  409. {
  410. "input_ids": tf.int32,
  411. "attention_mask": tf.int32,
  412. "token_type_ids": tf.int32,
  413. "feature_index": tf.int64,
  414. "qas_id": tf.string,
  415. },
  416. {
  417. "start_positions": tf.int64,
  418. "end_positions": tf.int64,
  419. "cls_index": tf.int64,
  420. "p_mask": tf.int32,
  421. "is_impossible": tf.int32,
  422. },
  423. )
  424. train_shapes = (
  425. {
  426. "input_ids": tf.TensorShape([None]),
  427. "attention_mask": tf.TensorShape([None]),
  428. "token_type_ids": tf.TensorShape([None]),
  429. "feature_index": tf.TensorShape([]),
  430. "qas_id": tf.TensorShape([]),
  431. },
  432. {
  433. "start_positions": tf.TensorShape([]),
  434. "end_positions": tf.TensorShape([]),
  435. "cls_index": tf.TensorShape([]),
  436. "p_mask": tf.TensorShape([None]),
  437. "is_impossible": tf.TensorShape([]),
  438. },
  439. )
  440. else:
  441. train_types = (
  442. {"input_ids": tf.int32, "attention_mask": tf.int32, "feature_index": tf.int64, "qas_id": tf.string},
  443. {
  444. "start_positions": tf.int64,
  445. "end_positions": tf.int64,
  446. "cls_index": tf.int64,
  447. "p_mask": tf.int32,
  448. "is_impossible": tf.int32,
  449. },
  450. )
  451. train_shapes = (
  452. {
  453. "input_ids": tf.TensorShape([None]),
  454. "attention_mask": tf.TensorShape([None]),
  455. "feature_index": tf.TensorShape([]),
  456. "qas_id": tf.TensorShape([]),
  457. },
  458. {
  459. "start_positions": tf.TensorShape([]),
  460. "end_positions": tf.TensorShape([]),
  461. "cls_index": tf.TensorShape([]),
  462. "p_mask": tf.TensorShape([None]),
  463. "is_impossible": tf.TensorShape([]),
  464. },
  465. )
  466. return tf.data.Dataset.from_generator(gen, train_types, train_shapes)
  467. else:
  468. return features
  469. class SquadProcessor(DataProcessor):
  470. """
  471. Processor for the SQuAD data set. overridden by SquadV1Processor and SquadV2Processor, used by the version 1.1 and
  472. version 2.0 of SQuAD, respectively.
  473. """
  474. train_file = None
  475. dev_file = None
  476. def _get_example_from_tensor_dict(self, tensor_dict, evaluate=False):
  477. if not evaluate:
  478. answer = tensor_dict["answers"]["text"][0].numpy().decode("utf-8")
  479. answer_start = tensor_dict["answers"]["answer_start"][0].numpy()
  480. answers = []
  481. else:
  482. answers = [
  483. {"answer_start": start.numpy(), "text": text.numpy().decode("utf-8")}
  484. for start, text in zip(tensor_dict["answers"]["answer_start"], tensor_dict["answers"]["text"])
  485. ]
  486. answer = None
  487. answer_start = None
  488. return SquadExample(
  489. qas_id=tensor_dict["id"].numpy().decode("utf-8"),
  490. question_text=tensor_dict["question"].numpy().decode("utf-8"),
  491. context_text=tensor_dict["context"].numpy().decode("utf-8"),
  492. answer_text=answer,
  493. start_position_character=answer_start,
  494. title=tensor_dict["title"].numpy().decode("utf-8"),
  495. answers=answers,
  496. )
  497. def get_examples_from_dataset(self, dataset, evaluate=False):
  498. """
  499. Creates a list of [`~data.processors.squad.SquadExample`] using a TFDS dataset.
  500. Args:
  501. dataset: The tfds dataset loaded from *tensorflow_datasets.load("squad")*
  502. evaluate: Boolean specifying if in evaluation mode or in training mode
  503. Returns:
  504. List of SquadExample
  505. Examples:
  506. ```python
  507. >>> import tensorflow_datasets as tfds
  508. >>> dataset = tfds.load("squad")
  509. >>> training_examples = get_examples_from_dataset(dataset, evaluate=False)
  510. >>> evaluation_examples = get_examples_from_dataset(dataset, evaluate=True)
  511. ```"""
  512. if evaluate:
  513. dataset = dataset["validation"]
  514. else:
  515. dataset = dataset["train"]
  516. examples = []
  517. for tensor_dict in tqdm(dataset):
  518. examples.append(self._get_example_from_tensor_dict(tensor_dict, evaluate=evaluate))
  519. return examples
  520. def get_train_examples(self, data_dir, filename=None):
  521. """
  522. Returns the training examples from the data directory.
  523. Args:
  524. data_dir: Directory containing the data files used for training and evaluating.
  525. filename: None by default, specify this if the training file has a different name than the original one
  526. which is `train-v1.1.json` and `train-v2.0.json` for squad versions 1.1 and 2.0 respectively.
  527. """
  528. if data_dir is None:
  529. data_dir = ""
  530. if self.train_file is None:
  531. raise ValueError("SquadProcessor should be instantiated via SquadV1Processor or SquadV2Processor")
  532. with open(
  533. os.path.join(data_dir, self.train_file if filename is None else filename), "r", encoding="utf-8"
  534. ) as reader:
  535. input_data = json.load(reader)["data"]
  536. return self._create_examples(input_data, "train")
  537. def get_dev_examples(self, data_dir, filename=None):
  538. """
  539. Returns the evaluation example from the data directory.
  540. Args:
  541. data_dir: Directory containing the data files used for training and evaluating.
  542. filename: None by default, specify this if the evaluation file has a different name than the original one
  543. which is `dev-v1.1.json` and `dev-v2.0.json` for squad versions 1.1 and 2.0 respectively.
  544. """
  545. if data_dir is None:
  546. data_dir = ""
  547. if self.dev_file is None:
  548. raise ValueError("SquadProcessor should be instantiated via SquadV1Processor or SquadV2Processor")
  549. with open(
  550. os.path.join(data_dir, self.dev_file if filename is None else filename), "r", encoding="utf-8"
  551. ) as reader:
  552. input_data = json.load(reader)["data"]
  553. return self._create_examples(input_data, "dev")
  554. def _create_examples(self, input_data, set_type):
  555. is_training = set_type == "train"
  556. examples = []
  557. for entry in tqdm(input_data):
  558. title = entry["title"]
  559. for paragraph in entry["paragraphs"]:
  560. context_text = paragraph["context"]
  561. for qa in paragraph["qas"]:
  562. qas_id = qa["id"]
  563. question_text = qa["question"]
  564. start_position_character = None
  565. answer_text = None
  566. answers = []
  567. is_impossible = qa.get("is_impossible", False)
  568. if not is_impossible:
  569. if is_training:
  570. answer = qa["answers"][0]
  571. answer_text = answer["text"]
  572. start_position_character = answer["answer_start"]
  573. else:
  574. answers = qa["answers"]
  575. example = SquadExample(
  576. qas_id=qas_id,
  577. question_text=question_text,
  578. context_text=context_text,
  579. answer_text=answer_text,
  580. start_position_character=start_position_character,
  581. title=title,
  582. is_impossible=is_impossible,
  583. answers=answers,
  584. )
  585. examples.append(example)
  586. return examples
  587. class SquadV1Processor(SquadProcessor):
  588. train_file = "train-v1.1.json"
  589. dev_file = "dev-v1.1.json"
  590. class SquadV2Processor(SquadProcessor):
  591. train_file = "train-v2.0.json"
  592. dev_file = "dev-v2.0.json"
  593. class SquadExample:
  594. """
  595. A single training/test example for the Squad dataset, as loaded from disk.
  596. Args:
  597. qas_id: The example's unique identifier
  598. question_text: The question string
  599. context_text: The context string
  600. answer_text: The answer string
  601. start_position_character: The character position of the start of the answer
  602. title: The title of the example
  603. answers: None by default, this is used during evaluation. Holds answers as well as their start positions.
  604. is_impossible: False by default, set to True if the example has no possible answer.
  605. """
  606. def __init__(
  607. self,
  608. qas_id,
  609. question_text,
  610. context_text,
  611. answer_text,
  612. start_position_character,
  613. title,
  614. answers=[],
  615. is_impossible=False,
  616. ):
  617. self.qas_id = qas_id
  618. self.question_text = question_text
  619. self.context_text = context_text
  620. self.answer_text = answer_text
  621. self.title = title
  622. self.is_impossible = is_impossible
  623. self.answers = answers
  624. self.start_position, self.end_position = 0, 0
  625. doc_tokens = []
  626. char_to_word_offset = []
  627. prev_is_whitespace = True
  628. # Split on whitespace so that different tokens may be attributed to their original position.
  629. for c in self.context_text:
  630. if _is_whitespace(c):
  631. prev_is_whitespace = True
  632. else:
  633. if prev_is_whitespace:
  634. doc_tokens.append(c)
  635. else:
  636. doc_tokens[-1] += c
  637. prev_is_whitespace = False
  638. char_to_word_offset.append(len(doc_tokens) - 1)
  639. self.doc_tokens = doc_tokens
  640. self.char_to_word_offset = char_to_word_offset
  641. # Start and end positions only has a value during evaluation.
  642. if start_position_character is not None and not is_impossible:
  643. self.start_position = char_to_word_offset[start_position_character]
  644. self.end_position = char_to_word_offset[
  645. min(start_position_character + len(answer_text) - 1, len(char_to_word_offset) - 1)
  646. ]
  647. class SquadFeatures:
  648. """
  649. Single squad example features to be fed to a model. Those features are model-specific and can be crafted from
  650. [`~data.processors.squad.SquadExample`] using the
  651. :method:*~transformers.data.processors.squad.squad_convert_examples_to_features* method.
  652. Args:
  653. input_ids: Indices of input sequence tokens in the vocabulary.
  654. attention_mask: Mask to avoid performing attention on padding token indices.
  655. token_type_ids: Segment token indices to indicate first and second portions of the inputs.
  656. cls_index: the index of the CLS token.
  657. p_mask: Mask identifying tokens that can be answers vs. tokens that cannot.
  658. Mask with 1 for tokens than cannot be in the answer and 0 for token that can be in an answer
  659. example_index: the index of the example
  660. unique_id: The unique Feature identifier
  661. paragraph_len: The length of the context
  662. token_is_max_context:
  663. List of booleans identifying which tokens have their maximum context in this feature object. If a token
  664. does not have their maximum context in this feature object, it means that another feature object has more
  665. information related to that token and should be prioritized over this feature for that token.
  666. tokens: list of tokens corresponding to the input ids
  667. token_to_orig_map: mapping between the tokens and the original text, needed in order to identify the answer.
  668. start_position: start of the answer token index
  669. end_position: end of the answer token index
  670. encoding: optionally store the BatchEncoding with the fast-tokenizer alignment methods.
  671. """
  672. def __init__(
  673. self,
  674. input_ids,
  675. attention_mask,
  676. token_type_ids,
  677. cls_index,
  678. p_mask,
  679. example_index,
  680. unique_id,
  681. paragraph_len,
  682. token_is_max_context,
  683. tokens,
  684. token_to_orig_map,
  685. start_position,
  686. end_position,
  687. is_impossible,
  688. qas_id: Optional[str] = None,
  689. encoding: Optional[BatchEncoding] = None,
  690. ):
  691. self.input_ids = input_ids
  692. self.attention_mask = attention_mask
  693. self.token_type_ids = token_type_ids
  694. self.cls_index = cls_index
  695. self.p_mask = p_mask
  696. self.example_index = example_index
  697. self.unique_id = unique_id
  698. self.paragraph_len = paragraph_len
  699. self.token_is_max_context = token_is_max_context
  700. self.tokens = tokens
  701. self.token_to_orig_map = token_to_orig_map
  702. self.start_position = start_position
  703. self.end_position = end_position
  704. self.is_impossible = is_impossible
  705. self.qas_id = qas_id
  706. self.encoding = encoding
  707. class SquadResult:
  708. """
  709. Constructs a SquadResult which can be used to evaluate a model's output on the SQuAD dataset.
  710. Args:
  711. unique_id: The unique identifier corresponding to that example.
  712. start_logits: The logits corresponding to the start of the answer
  713. end_logits: The logits corresponding to the end of the answer
  714. """
  715. def __init__(self, unique_id, start_logits, end_logits, start_top_index=None, end_top_index=None, cls_logits=None):
  716. self.start_logits = start_logits
  717. self.end_logits = end_logits
  718. self.unique_id = unique_id
  719. if start_top_index:
  720. self.start_top_index = start_top_index
  721. self.end_top_index = end_top_index
  722. self.cls_logits = cls_logits