"""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")