from datasets import load_dataset
from transformers import AutoTokenizer

# 设置模型 tokenizer（例如使用 GPT-2）
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")

# 加载 parquet 数据集（假设文件名为 data.parquet）
dataset = load_dataset("parquet", data_files="/workspace/0525_zyw/verl/counting/mk_data/counting_dataset_qwen25_max2048.parquet")

# 选择其中一个 split，例如 'train' 或默认的 'train'
data = dataset["train"]

# 记录超过 2048 tokens 的样本索引和内容
long_items = []

for idx, example in enumerate(data):
    prompt = example.get("prompt", "")
    tokens = tokenizer(prompt, truncation=False, return_tensors="pt")
    input_len = tokens.input_ids.shape[1]

    if input_len > 2048:
        long_items.append({"index": idx, "length": input_len, "prompt": prompt})

print(f"Found {len(long_items)} items with more than 2048 tokens.")

# 可选：保存结果到 JSON 文件
import json
with open("long_prompts.json", "w", encoding="utf-8") as f:
    json.dump(long_items, f, ensure_ascii=False, indent=2)