(self, code, hash, max_token_length)
| 56 | return {"results": results} |
| 57 | |
| 58 | def chunk_code(self, code, hash, max_token_length): |
| 59 | # Encode the entire code at once, ignoring special tokens |
| 60 | tokens_data = self.tokenizer.encode_plus(code, add_special_tokens=False, return_offsets_mapping=True) |
| 61 | tokens = tokens_data['input_ids'] |
| 62 | offsets = tokens_data['offset_mapping'] |
| 63 | |
| 64 | chunks = [] |
| 65 | current_chunk_start_index = 0 |
| 66 | total_tokens = len(tokens) |
| 67 | current_token_index = 0 |
| 68 | |
| 69 | # Distribute tokens to each chunk as evenly as possible with some padding |
| 70 | tokens_per_chunk = max_token_length |
| 71 | while current_token_index < total_tokens: |
| 72 | next_token_index = min(current_token_index + tokens_per_chunk, total_tokens) |
| 73 | |
| 74 | # Adjust the end index to not cut in the middle of a word |
| 75 | while next_token_index < total_tokens and offsets[next_token_index - 1][1] != offsets[next_token_index][0]: |
| 76 | next_token_index += 1 |
| 77 | |
| 78 | # Ensure not to exceed total tokens while correcting |
| 79 | next_token_index = min(next_token_index, total_tokens) |
| 80 | |
| 81 | chunk_tokens = tokens[current_chunk_start_index:next_token_index] |
| 82 | start_index = offsets[current_chunk_start_index][0] |
| 83 | end_index = offsets[next_token_index - 1][1] if next_token_index > 0 else 0 |
| 84 | |
| 85 | # Manually reconstruct the text to ensure spaces are correctly included |
| 86 | chunk_code = code[start_index:end_index] |
| 87 | |
| 88 | # Append the current chunk to the chunks list |
| 89 | chunks.append({ |
| 90 | "chunk_id": len(chunks), |
| 91 | "code": chunk_code, |
| 92 | "file_hash": hash, |
| 93 | "start_index": start_index, |
| 94 | "end_index": end_index |
| 95 | }) |
| 96 | |
| 97 | # Update for the next chunk |
| 98 | current_token_index = next_token_index |
| 99 | current_chunk_start_index = next_token_index |
| 100 | |
| 101 | return chunks |
| 102 | |
| 103 | |
| 104 |
no outgoing calls
no test coverage detected