#!/bin/bash
# Pure CPU preprocessing - no GPU, no DeepSpeed, just tokenize and cache
#SBATCH -A nvr_elm_llm  
#SBATCH -p cpu_short
#SBATCH -t 02:00:00
#SBATCH -N 1
#SBATCH -J preprocess_cpu
#SBATCH --cpus-per-task=64
#SBATCH -e workdir/preprocess_cpu_%j.err
#SBATCH -o workdir/preprocess_cpu_%j.out

set -eo pipefail

echo "[preprocess] CPU-only preprocessing - no GPU/DeepSpeed"
echo "[preprocess] Start at $(date)"

# activate environment
set +u
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda
set -u

cd /home/chengyuew/chengyuew/workspace/shiyi/LMFlow || cd /lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow

# Run pure Python preprocessing using LMFlow's actual tokenization logic
python3 << 'EOF'
import sys
import os
sys.path.insert(0, "/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src")
sys.path.insert(0, "/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow/src")
os.chdir("/home/chengyuew/chengyuew/workspace/shiyi/LMFlow")

print("[preprocess] Loading LMFlow components...")
from lmflow.datasets.dataset import Dataset
from lmflow.args import DatasetArguments, ModelArguments

print("[preprocess] Setting up arguments...")

# Use real DatasetArguments with all defaults
data_args = DatasetArguments(
    dataset_path="/home/chengyuew/chengyuew/workspace/LMFlow/data/nemotron_no_thinking_subset",
    conversation_template="qwen2_5_no_reasoning",
    block_size=2048,
    disable_group_texts=0,  # Match training script!
    preprocessing_num_workers=8,  # Reduced for stability
    overwrite_cache=False,
)

# Use real ModelArguments
model_args = ModelArguments(
    model_name_or_path="/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_DFM",
    tokenizer_name="Qwen/Qwen2.5-VL-3B-Instruct",
    trust_remote_code=True,
)

print(f"[preprocess] Loading tokenizer...")
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(
    model_args.tokenizer_name,
    trust_remote_code=model_args.trust_remote_code
)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
print(f"[preprocess] Tokenizer loaded")

print(f"[preprocess] Loading dataset from {data_args.dataset_path}")
dataset = Dataset(data_args=data_args, backend="huggingface")
print(f"[preprocess] Dataset loaded: {len(dataset)} examples")

print(f"[preprocess] Tokenizing with conversation template: {data_args.conversation_template}")
print(f"[preprocess] Using num_proc={data_args.preprocessing_num_workers}")

# Import the actual tokenization function from LMFlow
from lmflow.tokenization.hf_decoder_model import conversation_tokenize_function

# Create a wrapper that matches the expected signature
def tokenize_wrapper(examples):
    return conversation_tokenize_function(
        examples=examples,
        data_args=data_args,
        tokenizer=tokenizer,
        column_names=dataset.backend_dataset.column_names,
        conversation_template=data_args.conversation_template
    )

print(f"[preprocess] Starting tokenization (this will take ~30-50 minutes)...")

# Tokenize with caching - this is what training does
tokenized = dataset.backend_dataset.map(
    tokenize_wrapper,
    batched=True,
    num_proc=data_args.preprocessing_num_workers,
    load_from_cache_file=not data_args.overwrite_cache,
    desc="Tokenizing conversations",
    remove_columns=dataset.backend_dataset.column_names,
)

print(f"[preprocess] ✓✓✓ Tokenization complete! ✓✓✓")
print(f"[preprocess] Tokenized examples: {len(tokenized)}")
print(f"[preprocess] Cache generated in ~/.cache/huggingface/datasets/")
print(f"[preprocess] Done at {__import__('datetime').datetime.now()}")
EOF

echo "[preprocess] Script completed successfully at $(date)"
echo "[preprocess] Cache files:"
ls -lh ~/.cache/huggingface/datasets/ | tail -n 15

