import numpy as np
import torch
from transformers import AutoTokenizer

def read_bytes(path: str, bytes_start: int, num_bytes: int) -> bytes:
    with open(path, "rb") as f:
        f.seek(bytes_start)
        return f.read(num_bytes)

# 下面是读取整个文件的例子
path = "data/mmlu-sft/processed_data/nemotron5.5/tasks-data/n5.5-apr19-25-new/mmlu-sft/agieval_aqua_rat_aug/qwen3/data/0_00000.npy"  # 实际并不是 .npy 格式，可以命名成 .bin 更直观。 
dtype = np.uint32      # 根据对方说法，使用 uint32
bytes_per_element = 4  # uint32 是 4 个字节

# # 获取整个文件内容
# with open(path, "rb") as f:
#     file_bytes = f.read()

# # 转为 numpy array
# array = np.frombuffer(file_bytes, dtype=dtype)

# # 如果要转 PyTorch Tensor 且 int 类型：
# tensor = torch.tensor(array.astype(np.int64), dtype=torch.long)

# print(tensor)

index = 0               # 第几个块
chunk_size = 2048         # 每块 100 个元素（你可以换成实际值）
dtype = np.uint32
num_bytes = chunk_size * 4
offset = index * num_bytes

# 读取
buffer = read_bytes(path, offset, num_bytes)
array = np.frombuffer(buffer, dtype=dtype)
tensor = torch.tensor(array.astype(np.int64), dtype=torch.long)
print(tensor)

tokenizer = AutoTokenizer.from_pretrained("models/Qwen2.5-0.5B-Instruct")
print(tokenizer.decode(tensor))