| """PILE streaming + roleplay fine-tune data loading""" |
| def load_pile(seq_len, seed=42): |
| from datasets import load_dataset |
| ds = load_dataset("monology/pile-uncopyrighted", streaming=True, split="train") |
| def tok(ex): |
| t = ex.get("text", "") |
| if not t: t = "" |
| ts = list(t.encode("utf-8"))[:seq_len+1] |
| return {"tokens": ts} |
| ds = ds.map(tok).shuffle(seed=seed, buffer_size=10000) |
| return iter(ds) |
|
|
| def load_roleplay(seq_len, seed=123): |
| from datasets import load_dataset |
| for name in ["Kwangju/character_dialogue", "alexshcozh/roleplay-dialogues"]: |
| try: |
| ds = load_dataset(name, streaming=True, split="train") |
| def tok(ex): |
| t = ex.get("text", ex.get("dialogue", ex.get("input", ""))) |
| if not isinstance(t, str): t = str(t) if t else "" |
| ts = list(t.encode("utf-8"))[:seq_len+1] |
| return {"tokens": ts} |
| ds = ds.map(tok).shuffle(seed=seed, buffer_size=5000) |
| return iter(ds) |
| except: continue |
| return load_pile(seq_len, seed) |
|
|
| def make_batch(ds_it, bs, sl): |
| b = [] |
| for _ in range(bs): |
| try: |
| ex = next(ds_it); t = ex.get("tokens", [0]*(sl+1)) |
| if len(t) < sl+1: t += [0]*(sl+1-len(t)) |
| b.append([min(x,255) for x in t[:sl+1]]) |
| except StopIteration: b.append([0]*(sl+1)) |
| a = jnp.array(b, dtype=jnp.int32) |
| return a[:,:sl], a[:,1:] |
|
|
| print("Data loaders defined") |
|
|
