Chunk the large document between max and min context length of BERT-based model
(text, tokenizer, max_len=512, min_len=300, overlap_ratio=0.1)
| 64 | ) |
| 65 | |
| 66 | def chunk_document(text, tokenizer, max_len=512, min_len=300, overlap_ratio=0.1): |
| 67 | """Chunk the large document between max and min context length of BERT-based model""" |
| 68 | |
| 69 | tokenized = tokenizer(text, return_offsets_mapping=True, add_special_tokens=False) |
| 70 | input_ids = tokenized["input_ids"] |
| 71 | offsets = tokenized["offset_mapping"] |
| 72 | |
| 73 | total_tokens = len(input_ids) |
| 74 | overlap = int(max_len * overlap_ratio) # 10% overlap |
| 75 | chunks = [] |
| 76 | start = 0 |
| 77 | |
| 78 | while start <= total_tokens: |
| 79 | # Determine chunk length randomly (between min_len and max_len) - That's what our model is trained on |
| 80 | chunk_len = np.random.randint(min_len, max_len+1) |
| 81 | |
| 82 | chunk_end = min(start+chunk_len, total_tokens) |
| 83 | chunk_text = text[offsets[start][0]:offsets[chunk_end-1][1]] |
| 84 | chunks.append(chunk_text) |
| 85 | |
| 86 | if chunk_end == total_tokens: |
| 87 | break |
| 88 | if chunk_end - overlap>0: |
| 89 | start = chunk_end - overlap |
| 90 | else: |
| 91 | break |
| 92 | |
| 93 | return chunks |
| 94 | |
| 95 | |
| 96 | async def get_nosey_parker_results(doc): |