Get the Arrow file path (single file for all processes)

Forward pass that returns **fused log‑probs** as logits. We keep separate caches for each sub‑model.

Greedy decoding with **shared** stopping criteria. We keep two independent KV caches (one per sub‑model) and extend them step

Load FAISS index and optionally move to GPU

Save data for current step using streaming Arrow format

Set up the Arrow writer for streaming writes (only on main process)

Set up Arrow writer for streaming writes (main process only)

Build a FAISS index from an Arrow file containing keys and values. Args: dstore_path: Path to the Arrow file containing the data

Generate the path for the FAISS index file.

Compute kNN probability distribution following the reference implementation

Parse the dstore path to extract model_type, eval_subset, and dimension. Example: /fs-computility/plm/shared/jqcao/projects/neuralKN

Main processing loop

Select continuous chunks of data from non-overlapping regions of the dataset. Args: dataset: The dataset to sample from

Extract sparse representation of distribution with probs > threshold

Custom collate function that handles kNN data processing directly Args: batch: The batch of data knn