File size: 2,861 Bytes
97df2ae e26f4cc 97df2ae e26f4cc 97df2ae e26f4cc 97df2ae e26f4cc 97df2ae e26f4cc 97df2ae e26f4cc 97df2ae | 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 | #!/usr/bin/env python3
"""39_coletar_dados.py — Streaming CONTROLADO de dados reais (Wikipédia PT-BR).
Tetos duros (não satura o ambiente): max_artigos e max_bytes do config.
Salva shards JSONL locais (treino/validação 90/10) + treina o BPE do
KHTST_RUST (vocab-alvo do config) e salva em dados/bpe.json.
Nada além dos tetos é baixado; sem cache global do datasets."""
import json, os, sys, time
TETO_PADRAO_ARTIGOS = 1200
TETO_PADRAO_BYTES = 6_000_000
def main():
pasta = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
cfg = json.load(open(os.path.join(pasta, "configs", "base_rust.json")))
dados_cfg = cfg.get("dados", {})
max_artigos = int(dados_cfg.get("max_artigos", TETO_PADRAO_ARTIGOS))
max_bytes = int(dados_cfg.get("max_bytes", TETO_PADRAO_BYTES))
os.makedirs(os.path.join(pasta, "dados"), exist_ok=True)
from datasets import load_dataset
t0 = time.time()
ds = load_dataset(
"wikimedia/wikipedia", "20231101.pt", split="train", streaming=True,
)
textos, total_bytes = [], 0
for ex in ds:
t = (ex.get("text") or "").strip()
if len(t) < 200: # descarta esboços mínimos
continue
textos.append(t)
total_bytes += len(t.encode("utf-8"))
if len(textos) % 200 == 0:
print(f" ... {len(textos)} artigos, {total_bytes/1e6:.2f} MB", flush=True)
if len(textos) >= max_artigos or total_bytes >= max_bytes:
break
print(f"coletado: {len(textos)} artigos, {total_bytes/1e6:.2f} MB em {time.time()-t0:.1f}s")
# normalização leve: primeiros parágrafos relevantes (descarta infobox ruidosa)
def limpa(t: str) -> str:
partes = [p.strip() for p in t.split("\n") if len(p.strip()) > 80]
return " ".join(partes[:6])[:1800]
textos = [limpa(t) for t in textos]
n_val = max(1, len(textos) // 10)
val, treino = textos[:n_val], textos[n_val:]
with open(os.path.join(pasta, "dados", "shard_treino.jsonl"), "w") as f:
for t in treino:
f.write(json.dumps({"texto": t}, ensure_ascii=False) + "\n")
with open(os.path.join(pasta, "dados", "shard_validacao.jsonl"), "w") as f:
for t in val:
f.write(json.dumps({"texto": t}, ensure_ascii=False) + "\n")
print(f"shards: treino={len(treino)} val={len(val)}")
# BPE do KHTST_RUST
import khtst_rust as K
vocab_alvo = int(cfg["modelo"]["vocab"])
t0 = time.time()
bpe = K.Bpe.treinar(treino, vocab_alvo)
bpe.salvar(os.path.join(pasta, "dados", "bpe.json"))
print(f"BPE: vocab={bpe.n_vocab} (alvo {vocab_alvo}) em {time.time()-t0:.1f}s")
amostra = treino[0][:120]
ids = bpe.codificar(amostra)
rt = bpe.decodificar(ids)
print("roundtrip OK:", rt[:80] == amostra[:80].strip()[:80] or len(rt) > 50)
if __name__ == "__main__":
main()
|