Spaces:
Running
Running
| 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) | |