mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Ben Shoham Ofir
2a51bbd7cb
Optimize calibrate_draft_vocab to read only required lines when
calibrate_size is set
## What does this PR do?
**Type of change:** Performance improvement
**Overview:**
This PR optimizes the
[calibrate_draft_vocab.py](cci:7://file:///Users/obenshoham/PycharmProjects/TensorRT-Model-Optimizer/examples/speculative_decoding/scripts/calibrate_draft_vocab.py:0:0-0:0)
script to improve memory efficiency and I/O performance when using the
`--calibrate_size` parameter. Previously, the script would read all
lines from the data file into memory before slicing to the specified
`calibrate_size`, causing unnecessary resource usage for large datasets.
The optimization uses `itertools.islice` to read only the required
number of lines when `calibrate_size` is specified.
## Usage
The script usage remains unchanged. When using `--calibrate_size`, the
script now only reads the specified number of lines instead of loading
the entire dataset:
```bash
# Only reads first 1000 lines from the dataset (optimized)
python scripts/calibrate_draft_vocab.py \
--model meta-llama/Llama-3.2-1B-Instruct \
--data input_conversations/daring-anteater.jsonl \
--draft_vocab_size 32000 \
--calibrate_size 1000 \
--save_dir draft_vocab_cache
Signed-off-by: Ofir Ben Shoham <ofir_benshoham@intuit.com>