Set up Arrow writer for streaming writes (main process only)
(self)
| 109 | return cpu_index, gpu_index |
| 110 | |
| 111 | def _setup_arrow_writer(self): |
| 112 | """Set up Arrow writer for streaming writes (main process only)""" |
| 113 | if self.process_index == 0: |
| 114 | # Define schema for output Arrow file |
| 115 | fields = [ |
| 116 | pa.field('id_cnt', pa.int32()), |
| 117 | pa.field('token_id', pa.list_(pa.int32())), |
| 118 | pa.field('prob', pa.list_(pa.float16())), # Changed to float16 |
| 119 | pa.field('label', pa.int32()) |
| 120 | ] |
| 121 | schema = pa.schema(fields) |
| 122 | |
| 123 | # Create output directory if needed |
| 124 | Path(self.output_path).parent.mkdir(parents=True, exist_ok=True) |
| 125 | |
| 126 | # Create Arrow file and writer |
| 127 | self.arrow_file = pa.OSFile(self.output_path, 'wb') |
| 128 | self.arrow_writer = pa.ipc.new_stream(self.arrow_file, schema) |
| 129 | |
| 130 | def get_knns(self, queries, ignore_first=False): |
| 131 | if not self.knn_gpu: |