Custom Code for the Just Read Twice Project
(self,
doc, doc_id, tokenizer = None, context_length = 1000, answer_length = 50, sequence_length = 2048, context_key = "context", cutting_context = False, answer_key = ["value"], id_start=0, skip_check_answer=False,
)
| 350 | |
| 351 | |
| 352 | def truncate_context(self, |
| 353 | doc, doc_id, tokenizer = None, context_length = 1000, answer_length = 50, sequence_length = 2048, context_key = "context", cutting_context = False, answer_key = ["value"], id_start=0, skip_check_answer=False, |
| 354 | ): |
| 355 | """ |
| 356 | Custom Code for the Just Read Twice Project |
| 357 | """ |
| 358 | |
| 359 | instances = [] |
| 360 | new_doc_set = {} |
| 361 | desired_length = context_length - answer_length # buffer for the answer |
| 362 | context = doc[context_key] |
| 363 | context = context.strip() |
| 364 | |
| 365 | doc_tokens = tokenizer.batch_encode_plus([context], return_tensors="pt", padding=True, truncation=False)['input_ids'][0] |
| 366 | # doc_tokens = tokenizer.batch_encode_plus([context], return_tensors="pt", padding=True, truncation=True, max_length=sequence_length)['input_ids'][0] |
| 367 | |
| 368 | # Find where the answer is in the document |
| 369 | answer_pos = -1 |
| 370 | for key in answer_key: |
| 371 | answer = doc[key] |
| 372 | if type(answer) == list: answer = answer[0] |
| 373 | if(answer == "" or len(answer) <= 1): return instances, new_doc_set |
| 374 | answer_pattern = re.compile(re.escape(answer), re.IGNORECASE) |
| 375 | if answer_match := answer_pattern.search(context): |
| 376 | if answer_pos == -1 or answer_pos > answer_match.start(): |
| 377 | answer_pos = answer_match.start() |
| 378 | if answer_pos < 0 and 'alt_answers' in doc: |
| 379 | answers = doc['alt_answers'] |
| 380 | for answer in answers: |
| 381 | answer_pattern = re.compile(re.escape(answer), re.IGNORECASE) |
| 382 | if answer_match := answer_pattern.search(context): |
| 383 | if answer_pos == -1 or answer_pos > answer_match.start(): |
| 384 | answer_pos = answer_match.start() |
| 385 | if answer_pos >= 0: break |
| 386 | if len(answer_key) == 0: # E.g. summarization datasets |
| 387 | answer_pos = 0 |
| 388 | if skip_check_answer: |
| 389 | assert answer_pos >= 0, f"Did not find answer in context: {doc_id}" |
| 390 | |
| 391 | # Convert the answer_pos to a token value |
| 392 | context_short = context[:answer_pos] |
| 393 | context_short_tokens = tokenizer.batch_encode_plus( |
| 394 | [context_short], return_tensors="pt", |
| 395 | )['input_ids'][0] |
| 396 | answer_tok_pos = len(context_short_tokens) |
| 397 | |
| 398 | if answer_pos == -1 and 'tok_pos' in doc and len(context_short_tokens) > desired_length: |
| 399 | answer_tok_pos = doc['tok_pos'] |
| 400 | |
| 401 | # Pick new bounds |
| 402 | half_length = desired_length // 2 |
| 403 | start = max(0, answer_tok_pos - half_length) |
| 404 | completed_length = answer_tok_pos - start |
| 405 | remaining_length = desired_length - completed_length |
| 406 | end = min(len(doc_tokens), answer_tok_pos + remaining_length) |
| 407 | subset_tokens = doc_tokens[start:end] |
| 408 | short_context = tokenizer.decode(subset_tokens, skip_special_tokens=True) |
| 409 |
no test coverage detected