82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from dataset.lm_dataset import SFTDataset, post_processing_chat, pre_processing_chat
|
|
|
|
|
|
def decode_one(tokenizer, token_id):
|
|
if token_id == -100:
|
|
return "<ignored>"
|
|
return tokenizer.decode([int(token_id)], skip_special_tokens=False).replace("\n", "\\n")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Inspect MiniMind SFTDataset prompt and labels.")
|
|
parser.add_argument("--data_path", default=str(ROOT / "course/labs/tiny_sft.jsonl"))
|
|
parser.add_argument("--tokenizer_path", default=str(ROOT / "model"))
|
|
parser.add_argument("--index", type=int, default=0)
|
|
parser.add_argument("--max_length", type=int, default=96)
|
|
parser.add_argument("--show_rows", type=int, default=80)
|
|
args = parser.parse_args()
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
|
ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_length)
|
|
|
|
with open(args.data_path, "r", encoding="utf-8") as f:
|
|
raw_sample = json.loads(f.readlines()[args.index])
|
|
|
|
conversations = pre_processing_chat(raw_sample["conversations"], add_system_ratio=0.0)
|
|
prompt = ds.create_chat_prompt(conversations)
|
|
prompt = post_processing_chat(prompt, empty_think_ratio=1.0)
|
|
input_ids = tokenizer(prompt).input_ids[: args.max_length]
|
|
input_ids += [tokenizer.pad_token_id] * (args.max_length - len(input_ids))
|
|
labels = ds.generate_labels(input_ids)
|
|
input_ids = torch.tensor(input_ids, dtype=torch.long)
|
|
labels = torch.tensor(labels, dtype=torch.long)
|
|
|
|
print("[Raw conversations]")
|
|
print(json.dumps(raw_sample["conversations"], ensure_ascii=False, indent=2))
|
|
print()
|
|
|
|
print("[Chat prompt]")
|
|
print(prompt)
|
|
print()
|
|
|
|
print("[Shapes]")
|
|
print(f"input_ids.shape={tuple(input_ids.shape)}")
|
|
print(f"labels.shape={tuple(labels.shape)}")
|
|
print(f"non_ignored_labels={(labels != -100).sum().item()}")
|
|
print()
|
|
|
|
print("[Assistant span markers]")
|
|
print(f"bos_id={ds.bos_id} -> {tokenizer.convert_ids_to_tokens(ds.bos_id)}")
|
|
print(f"eos_id={ds.eos_id} -> {tokenizer.convert_ids_to_tokens(ds.eos_id)}")
|
|
print()
|
|
|
|
print("[Token / label table]")
|
|
print("idx | input_id | token | label | label_token")
|
|
print("-" * 72)
|
|
rows = min(args.show_rows, args.max_length)
|
|
for i in range(rows):
|
|
tid = int(input_ids[i])
|
|
label = int(labels[i])
|
|
token = tokenizer.decode([tid], skip_special_tokens=False).replace("\n", "\\n")
|
|
label_token = decode_one(tokenizer, label)
|
|
print(f"{i:03d} | {tid:8d} | {token!r:16s} | {label:6d} | {label_token!r}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
torch.set_printoptions(linewidth=160)
|
|
main()
|