import pickle
import os

BASE_DIR = "/home/mkd-s003/AI/finetune"
TRAIN_PKL = os.path.join(BASE_DIR, "token_cache_train.pkl")
VAL_PKL = os.path.join(BASE_DIR, "token_cache_val.pkl")

def count_tokens(pkl_path, label):
    total_tokens = 0
    total_sequences = 0
    with open(pkl_path, "rb") as f:
        while True:
            try:
                seq = pickle.load(f)
                total_tokens += len(seq)
                total_sequences += 1
            except EOFError:
                break
    print(f"{label}: {total_tokens:,} tokens in {total_sequences:,} sequences")
    return total_tokens

train_total = count_tokens(TRAIN_PKL, "Train")
val_total = count_tokens(VAL_PKL, "Validation")

print("\n========== SUMMARY ==========")
print(f"Total tokens (Train + Val): {train_total + val_total:,}")
print("=============================")