Spaces:
Running
Running
File size: 3,978 Bytes
8fa3dd6 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | from __future__ import annotations
import json
from pathlib import Path
from typing import Iterable, Iterator, List, Sequence
def read_text_records(paths: Sequence[str]) -> Iterator[str]:
"""Yield text records from txt/md/json/jsonl files or directories."""
files: List[Path] = []
for item in paths:
p = Path(item)
if p.is_dir():
files.extend([x for x in sorted(p.rglob("*")) if x.is_file()])
elif p.is_file():
files.append(p)
else:
raise FileNotFoundError(item)
for p in files:
suffix = p.suffix.lower()
if suffix in {".txt", ".md"}:
text = p.read_text(encoding="utf-8", errors="ignore").strip()
if text:
yield text
elif suffix == ".jsonl":
with p.open("r", encoding="utf-8", errors="ignore") as f:
for line in f:
line = line.strip()
if not line:
continue
obj = json.loads(line)
yield extract_text_from_json(obj)
elif suffix == ".json":
obj = json.loads(p.read_text(encoding="utf-8", errors="ignore"))
if isinstance(obj, list):
for row in obj:
yield extract_text_from_json(row)
else:
yield extract_text_from_json(obj)
def extract_text_from_json(obj) -> str:
if isinstance(obj, str):
return obj
if isinstance(obj, dict):
if "text" in obj:
return str(obj["text"])
if "content" in obj:
return str(obj["content"])
if "messages" in obj:
return format_messages(obj["messages"])
# Conservative fallback: concatenate scalar values.
parts = []
for value in obj.values():
if isinstance(value, (str, int, float)):
parts.append(str(value))
return "\n".join(parts)
return str(obj)
def format_messages(messages) -> str:
parts = []
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
parts.append(f"<|{role}|>\n{content}\n<|end|>")
return "\n".join(parts)
class PackedTokenDataset:
"""Simple in-memory packed causal-LM dataset.
For serious training, replace this with streaming shards / mmap arrays. This class is intentionally
small and readable for the first Ares smoke tests.
"""
def __init__(self, token_ids: Sequence[int], seq_len: int):
if len(token_ids) < seq_len + 1:
raise ValueError("Not enough tokens for one training example")
self.token_ids = list(map(int, token_ids))
self.seq_len = int(seq_len)
self.n = (len(self.token_ids) - 1) // self.seq_len
def __len__(self) -> int:
return self.n
def __getitem__(self, idx: int):
import torch
start = (idx % self.n) * self.seq_len
chunk = self.token_ids[start : start + self.seq_len + 1]
x = torch.tensor(chunk[:-1], dtype=torch.long)
y = torch.tensor(chunk[1:], dtype=torch.long)
return x, y
def encode_corpus(tokenizer_path: str, text_paths: Sequence[str]) -> List[int]:
try:
from tokenizers import Tokenizer
except ImportError as exc:
raise SystemExit("Install tokenizers first: pip install tokenizers") from exc
tok = Tokenizer.from_file(tokenizer_path)
ids: List[int] = []
for record in read_text_records(text_paths):
enc = tok.encode(record)
ids.extend(enc.ids)
return ids
def make_dataloader(tokenizer_path: str, text_paths: Sequence[str], seq_len: int, batch_size: int, shuffle: bool = True):
import torch
from torch.utils.data import DataLoader
ids = encode_corpus(tokenizer_path, text_paths)
dataset = PackedTokenDataset(ids, seq_len)
return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=True)
|