#!/bin/bash

mkdir -p $MTP_ROOT/outputs/results

device=cuda
num_tokens=1024
default="$MTP_ROOT/outputs/models/"
subsample_prompts=250
prompt_source="tulu-valid"
directory="${1:-$default}"
filetag="${2:-unspecified}"
prompt_index="${3:-0}"
step="${4:-900}"

run_id="throughput-evabyte-sampling-${filetag}-${num_tokens}-${subsample_prompts}"

export GPUS=1
export GPU=1

for file in $(find "$directory" -type f -name "model@${step}.pt" -not -path '*/.*');
do
  echo -e "Processing $file..."
  python -m mtp.generate \
    --device $device \
    --mode mtp \
    --use-cache \
    --speculative \
    --random-seed 42 \
    --prompt-subset-index $prompt_index \
    --subsample-prompts $subsample_prompts \
    --prompt-source $prompt_source \
    --num-tokens $num_tokens \
    --task chat \
    --checkpoint $file \
    --run-id $run_id \
    | tail -n 1 >> "$MTP_ROOT/outputs/results/${run_id}.jsonl"
done
