Spaces:
Sleeping
Sleeping
File size: 3,655 Bytes
8125804 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import json
from pathlib import Path
import torch
from src.data.the_stack import _normalize_repo_name
from src.data.the_stack_bpe import BPETokenDataset, _train_tokenizer
def _load_tinystories_text(
data_dir: str,
repo_id: str,
filename: str,
target_bytes: int,
) -> str:
from huggingface_hub import hf_hub_download
Path(data_dir).mkdir(parents=True, exist_ok=True)
cache_name = f"{_normalize_repo_name(repo_id)}_{filename.replace('.', '_')}_{target_bytes}.txt"
cache_path = Path(data_dir) / cache_name
if cache_path.exists():
return cache_path.read_text(encoding="utf-8")
src_path = hf_hub_download(repo_id, filename, repo_type="dataset")
with open(src_path, "r", encoding="utf-8") as f:
text = f.read(target_bytes)
cache_path.write_text(text, encoding="utf-8")
return text
def load_tinystories_bpe(
seq_len: int = 256,
device: str = "cpu",
data_dir: str = "data_cache",
repo_id: str = "roneneldan/TinyStories",
train_filename: str = "TinyStories-train.txt",
val_filename: str = "TinyStories-valid.txt",
target_bytes: int = 16_000_000,
vocab_size: int = 4096,
) -> tuple[BPETokenDataset, BPETokenDataset]:
Path(data_dir).mkdir(parents=True, exist_ok=True)
prefix = f"{_normalize_repo_name(repo_id)}_tinystories_{target_bytes}_bpe{vocab_size}"
tokenizer_path = Path(data_dir) / f"{prefix}_tokenizer.json"
train_ids_path = Path(data_dir) / f"{prefix}_train_ids.pt"
val_ids_path = Path(data_dir) / f"{prefix}_val_ids.pt"
meta_path = Path(data_dir) / f"{prefix}_meta.json"
if tokenizer_path.exists() and train_ids_path.exists() and val_ids_path.exists() and meta_path.exists():
train_ids = torch.load(train_ids_path, map_location="cpu")
val_ids = torch.load(val_ids_path, map_location="cpu")
meta = json.loads(meta_path.read_text(encoding="utf-8"))
actual_vocab_size = int(meta["vocab_size"])
else:
train_text = _load_tinystories_text(
data_dir=data_dir,
repo_id=repo_id,
filename=train_filename,
target_bytes=target_bytes,
)
val_text = _load_tinystories_text(
data_dir=data_dir,
repo_id=repo_id,
filename=val_filename,
target_bytes=max(1_000_000, target_bytes // 8),
)
tokenizer = _train_tokenizer(text=train_text, vocab_size=vocab_size)
train_ids = torch.tensor(tokenizer.encode(train_text).ids, dtype=torch.long)
val_ids = torch.tensor(tokenizer.encode(val_text).ids, dtype=torch.long)
actual_vocab_size = tokenizer.get_vocab_size()
tokenizer.save(str(tokenizer_path))
torch.save(train_ids, train_ids_path)
torch.save(val_ids, val_ids_path)
meta_path.write_text(
json.dumps(
{
"repo_id": repo_id,
"target_bytes": target_bytes,
"vocab_size": actual_vocab_size,
"train_token_count": int(train_ids.numel()),
"val_token_count": int(val_ids.numel()),
},
indent=2,
),
encoding="utf-8",
)
train = BPETokenDataset(
token_ids=train_ids,
vocab_size=actual_vocab_size,
split="train",
seq_len=seq_len,
device=device,
split_data=False,
)
val = BPETokenDataset(
token_ids=val_ids,
vocab_size=actual_vocab_size,
split="val",
seq_len=seq_len,
device=device,
split_data=False,
)
return train, val
|