pytensorforge 0.1.0__py3-none-any.whl

src/data/shard_builder.py

@@ -0,0 +1,119 @@
+import hashlib
+import json
+import os
+import time
+from multiprocessing import get_context
+
+from src.data.document_stream import DocumentReader
+from src.data.shard_writer import SHARD_FORMAT_VERSION, ShardWriter, dtype_for_vocab
+from src.tokenization.registry import load_tokenizer
+
+MANIFEST_NAME = "manifest.json"
+
+
+def corpus_version(files):
+    digest = hashlib.sha256()
+
+    for path in files:
+        stat = os.stat(path)
+        digest.update(f"{path}:{stat.st_size}:{int(stat.st_mtime)}".encode())
+
+    return digest.hexdigest()
+
+
+def _partition(files, workers):
+    workers = max(1, min(workers, len(files)))
+    size, extra = divmod(len(files), workers)
+    parts = []
+    start = 0
+
+    for i in range(workers):
+        end = start + size + (1 if i < extra else 0)
+        parts.append(files[start:end])
+        start = end
+
+    return parts
+
+
+def _worker(args):
+    worker_id, files, tokenizer_path, output_dir, shard_tokens, insert_eos, text_field, read_buffer_size = args
+
+    tokenizer = load_tokenizer(tokenizer_path)
+    writer = ShardWriter(
+        output_dir,
+        f"shard-w{worker_id:03d}",
+        shard_tokens,
+        dtype_for_vocab(tokenizer.vocab_size),
+    )
+
+    reader = DocumentReader(files, read_buffer_size=read_buffer_size, text_field=text_field)
+    eos = tokenizer.eos_id
+
+    for _, _, text, doc_end in reader.iter_documents():
+        if text:
+            writer.write(tokenizer.encode(text))
+
+        if doc_end and insert_eos and eos is not None:
+            writer.write([eos])
+
+    return writer.close()
+
+
+def build_shards(
+    corpus,
+    tokenizer_path,
+    output_dir,
+    shard_tokens=50_000_000,
+    workers=1,
+    insert_eos=True,
+    text_field="text",
+    read_buffer_size=1 << 20,
+):
+    if not corpus.files:
+        raise ValueError("corpus contains no files")
+
+    tokenizer = load_tokenizer(tokenizer_path)
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    parts = _partition(corpus.files, workers)
+
+    jobs = [
+        (i, part, tokenizer_path, output_dir, shard_tokens, insert_eos, text_field, read_buffer_size)
+        for i, part in enumerate(parts)
+    ]
+
+    started = time.time()
+
+    if len(jobs) == 1:
+        results = [_worker(jobs[0])]
+    else:
+        with get_context("spawn").Pool(len(jobs)) as pool:
+            results = pool.map(_worker, jobs)
+
+    shards = [shard for result in results for shard in result]
+
+    manifest = {
+        "format_version": SHARD_FORMAT_VERSION,
+        "tokenizer_identity": tokenizer.identity,
+        "vocab_size": tokenizer.vocab_size,
+        "dtype": dtype_for_vocab(tokenizer.vocab_size).name,
+        "insert_eos": insert_eos,
+        "corpus_version": corpus_version(corpus.files),
+        "source_files": list(corpus.files),
+        "total_tokens": sum(s["token_count"] for s in shards),
+        "shards": shards,
+        "created_at": time.time(),
+        "build_seconds": time.time() - started,
+    }
+
+    tmp = os.path.join(output_dir, MANIFEST_NAME + ".tmp")
+
+    with open(tmp, "w") as f:
+        json.dump(manifest, f, indent=2)
+        f.flush()
+        os.fsync(f.fileno())
+
+    os.replace(tmp, os.path.join(output_dir, MANIFEST_NAME))
+
+    return manifest

src/data/shard_writer.py

@@ -0,0 +1,81 @@
+import hashlib
+import os
+
+import numpy as np
+
+SHARD_FORMAT_VERSION = 1
+
+
+def dtype_for_vocab(vocab_size):
+    return np.dtype(np.uint16) if vocab_size <= 65535 else np.dtype(np.uint32)
+
+
+def file_sha256(path, block_size=1 << 20):
+    digest = hashlib.sha256()
+
+    with open(path, "rb") as f:
+        while True:
+            block = f.read(block_size)
+            if not block:
+                break
+            digest.update(block)
+
+    return digest.hexdigest()
+
+
+class ShardWriter:
+
+    def __init__(self, output_dir, prefix, shard_tokens, dtype):
+        self.output_dir = output_dir
+        self.prefix = prefix
+        self.shard_tokens = shard_tokens
+        self.dtype = np.dtype(dtype)
+
+        self._pending = []
+        self._pending_count = 0
+        self._shard_index = 0
+        self.shards = []
+
+        os.makedirs(output_dir, exist_ok=True)
+
+    def write(self, token_ids):
+        self._pending.append(np.asarray(token_ids, dtype=self.dtype))
+        self._pending_count += len(token_ids)
+
+        while self._pending_count >= self.shard_tokens:
+            self._flush(self.shard_tokens)
+
+    def _flush(self, count):
+        merged = np.concatenate(self._pending) if len(self._pending) > 1 else self._pending[0]
+
+        chunk = merged[:count]
+        rest = merged[count:]
+
+        self._pending = [rest] if rest.size else []
+        self._pending_count = int(rest.size)
+
+        name = f"{self.prefix}-{self._shard_index:05d}.bin"
+        final_path = os.path.join(self.output_dir, name)
+        tmp_path = final_path + ".tmp"
+
+        with open(tmp_path, "wb") as f:
+            chunk.tofile(f)
+            f.flush()
+            os.fsync(f.fileno())
+
+        os.replace(tmp_path, final_path)
+
+        self.shards.append({
+            "file": name,
+            "token_count": int(chunk.size),
+            "dtype": self.dtype.name,
+            "checksum": file_sha256(final_path),
+        })
+
+        self._shard_index += 1
+
+    def close(self):
+        if self._pending_count > 0:
+            self._flush(self._pending_count)
+
+        return self.shards

src/data/sharded_dataset.py

@@ -0,0 +1,112 @@
+import json
+import os
+
+import numpy as np
+
+from src.data.shard_builder import MANIFEST_NAME
+from src.tokenization.base import identity_matches
+
+
+def is_shard_dir(path):
+    return isinstance(path, str) and os.path.isfile(os.path.join(path, MANIFEST_NAME))
+
+
+class ShardIndex:
+
+    def __init__(self, directory):
+        self.directory = directory
+
+        with open(os.path.join(directory, MANIFEST_NAME)) as f:
+            self.manifest = json.load(f)
+
+        self.files = [os.path.join(directory, s["file"]) for s in self.manifest["shards"]]
+
+    def __len__(self):
+        return len(self.files)
+
+    def state_dict(self):
+        return {
+            "corpus_version": self.manifest["corpus_version"],
+            "total_tokens": self.manifest["total_tokens"],
+            "shards": [s["file"] for s in self.manifest["shards"]],
+        }
+
+    def load_state_dict(self, state):
+        if state.get("shards") != [s["file"] for s in self.manifest["shards"]] or \
+                state.get("corpus_version") != self.manifest["corpus_version"]:
+            raise ValueError("token shard set does not match the checkpoint; resuming would corrupt training")
+
+
+class ShardedTokenDataset:
+
+    def __init__(self, directory, sequence_length, tokenizer=None):
+        self.corpus = ShardIndex(directory)
+        self.sequence_length = sequence_length
+        self.dtype = np.dtype(self.corpus.manifest["dtype"])
+
+        if tokenizer is not None and not identity_matches(self.corpus.manifest["tokenizer_identity"], tokenizer.identity):
+            raise ValueError("token shards were built with a different tokenizer than the one supplied")
+
+        self.shard_index = 0
+        self.token_offset = 0
+        self._carry = np.empty(0, dtype=np.int64)
+
+    def batches(self, batch_size):
+        needed = self.sequence_length + 1
+        meta = self.corpus.manifest["shards"]
+        read_tokens = max(needed * batch_size * 8, 1 << 16)
+
+        while True:
+            examples = []
+
+            while len(examples) < batch_size:
+                while self._carry.size < needed and self.shard_index < len(meta):
+                    path = self.corpus.files[self.shard_index]
+                    count = meta[self.shard_index]["token_count"]
+
+                    if self.token_offset >= count:
+                        self.shard_index += 1
+                        self.token_offset = 0
+                        continue
+
+                    take = min(read_tokens, count - self.token_offset)
+                    chunk = np.fromfile(
+                        path,
+                        dtype=self.dtype,
+                        count=take,
+                        offset=self.token_offset * self.dtype.itemsize,
+                    ).astype(np.int64)
+
+                    self.token_offset += take
+                    self._carry = np.concatenate([self._carry, chunk]) if self._carry.size else chunk
+
+                if self._carry.size < needed:
+                    break
+
+                examples.append(self._carry[:needed])
+                self._carry = self._carry[needed:]
+
+            if not examples:
+                return
+
+            arr = np.stack(examples)
+
+            yield arr[:, :-1], arr[:, 1:]
+
+    def start_new_epoch(self):
+        self.shard_index = 0
+        self.token_offset = 0
+        self._carry = np.empty(0, dtype=np.int64)
+
+    def state_dict(self):
+        return {
+            "kind": "shards",
+            "shard_index": self.shard_index,
+            "token_offset": self.token_offset,
+            "carry": self._carry.tolist(),
+        }
+
+    def load_state_dict(self, data):
+        self.shard_index = data["shard_index"]
+        self.token_offset = data["token_offset"]
+        self._carry = np.asarray(data["carry"], dtype=np.int64)

src/data/streaming_dataset.py

@@ -0,0 +1,132 @@
+from collections import deque
+
+import numpy as np
+
+from src.data.document_stream import DocumentReader
+from src.data.parallel_encode import ParallelEncoder, encode_serial
+
+
+class DatasetState:
+
+    def __init__(self, file_index=0, byte_offset=0, token_buffer=None):
+        self.file_index = file_index
+        self.byte_offset = byte_offset
+        self.token_buffer = list(token_buffer) if token_buffer else []
+
+    def to_dict(self):
+        return {
+            "kind": "text",
+            "file_index": self.file_index,
+            "byte_offset": self.byte_offset,
+            "token_buffer": list(self.token_buffer),
+        }
+
+    @classmethod
+    def from_dict(cls, data):
+        return cls(
+            file_index=data.get("file_index", 0),
+            byte_offset=data.get("byte_offset", 0),
+            token_buffer=data.get("token_buffer", []),
+        )
+
+
+class StreamingTextDataset:
+
+    def __init__(
+        self,
+        corpus,
+        tokenizer,
+        sequence_length,
+        read_buffer_size=1 << 20,
+        insert_eos=True,
+        text_field="text",
+        state=None,
+        workers=1,
+    ):
+        self.corpus = corpus
+        self.tokenizer = tokenizer
+        self.sequence_length = sequence_length
+        self.insert_eos = insert_eos
+        self.workers = max(1, int(workers))
+
+        self.reader = DocumentReader(
+            corpus.files,
+            read_buffer_size=read_buffer_size,
+            text_field=text_field,
+        )
+
+        self.state = state or DatasetState()
+        self.token_buffer = deque(self.state.token_buffer)
+
+    def _encoded_documents(self, doc_iter):
+        if self.workers > 1:
+            return ParallelEncoder(self.tokenizer, self.workers).encode(doc_iter)
+
+        return encode_serial(self.tokenizer, doc_iter)
+
+    def _fill_buffer(self, doc_iter, min_tokens):
+        while len(self.token_buffer) < min_tokens:
+            try:
+                file_index, offset, ids, doc_end = next(doc_iter)
+            except StopIteration:
+                return False
+
+            self.state.file_index = file_index
+            self.state.byte_offset = offset
+
+            if len(ids):
+                self.token_buffer.extend(ids.tolist() if isinstance(ids, np.ndarray) else ids)
+
+            if doc_end and self.insert_eos and self.tokenizer.eos_id is not None:
+                self.token_buffer.append(self.tokenizer.eos_id)
+
+        return True
+
+    def batches(self, batch_size):
+        needed = self.sequence_length + 1
+        doc_iter = self._encoded_documents(self.reader.iter_documents(
+            self.state.file_index,
+            self.state.byte_offset,
+        ))
+
+        try:
+            yield from self._batches(doc_iter, batch_size, needed)
+        finally:
+            doc_iter.close()
+
+    def _batches(self, doc_iter, batch_size, needed):
+        while True:
+            examples = []
+
+            for _ in range(batch_size):
+                has_more = self._fill_buffer(doc_iter, needed)
+
+                if len(self.token_buffer) < needed:
+                    break
+
+                block = [self.token_buffer.popleft() for _ in range(needed)]
+                examples.append(block)
+
+                if not has_more and len(self.token_buffer) < needed:
+                    break
+
+            if not examples:
+                return
+
+            arr = np.array(examples, dtype=np.int64)
+
+            self.state.token_buffer = list(self.token_buffer)
+
+            yield arr[:, :-1], arr[:, 1:]
+
+    def start_new_epoch(self):
+        self.state = DatasetState()
+        self.token_buffer = deque()
+
+    def state_dict(self):
+        self.state.token_buffer = list(self.token_buffer)
+        return self.state.to_dict()
+
+    def load_state_dict(self, data):
+        self.state = DatasetState.from_dict(data)
+        self.token_buffer = deque(self.state.token_buffer)

src/data/validation.py

@@ -0,0 +1,212 @@
+import json
+import os
+
+from src.data.document_stream import JSON_EXTENSIONS
+from src.data.shard_builder import MANIFEST_NAME, corpus_version
+from src.data.shard_writer import file_sha256
+from src.tokenization.base import identity_matches
+
+
+class ValidationReport:
+
+    def __init__(self):
+        self.errors = []
+        self.warnings = []
+        self.stats = {}
+
+    @property
+    def ok(self):
+        return not self.errors
+
+    def to_dict(self):
+        return {"ok": self.ok, "errors": self.errors, "warnings": self.warnings, "stats": self.stats}
+
+
+def validate_corpus(corpus, text_field="text", sample_bytes=4 << 20, tokenizer=None, sample_docs=200):
+    report = ValidationReport()
+
+    if not corpus.files:
+        report.errors.append("corpus contains no supported files")
+        return report
+
+    total_bytes = 0
+    malformed = 0
+    bad_encoding = 0
+    empty_docs = 0
+    checked_docs = 0
+    unk_tokens = 0
+    sampled_tokens = 0
+    sampled_chars = 0
+
+    for path in corpus.files:
+        if not os.path.isfile(path):
+            report.errors.append(f"missing file: {path}")
+            continue
+
+        if not os.access(path, os.R_OK):
+            report.errors.append(f"unreadable file: {path}")
+            continue
+
+        total_bytes += os.path.getsize(path)
+
+    report.stats["files"] = len(corpus.files)
+    report.stats["total_bytes"] = total_bytes
+
+    remaining = sample_bytes
+
+    for path in corpus.files:
+        if remaining <= 0 or checked_docs >= sample_docs * 10:
+            break
+
+        if not os.path.isfile(path):
+            continue
+
+        ext = os.path.splitext(path)[1].lower()
+
+        with open(path, "rb") as f:
+            if ext in JSON_EXTENSIONS:
+                for line in f:
+                    remaining -= len(line)
+
+                    if not line.strip():
+                        continue
+
+                    checked_docs += 1
+
+                    try:
+                        text = line.decode("utf-8")
+                    except UnicodeDecodeError:
+                        bad_encoding += 1
+                        continue
+
+                    try:
+                        obj = json.loads(text)
+                    except json.JSONDecodeError:
+                        malformed += 1
+                        continue
+
+                    value = obj.get(text_field) if isinstance(obj, dict) else None
+
+                    if not isinstance(value, str):
+                        malformed += 1
+                    elif not value.strip():
+                        empty_docs += 1
+                    elif tokenizer is not None and sampled_tokens < 200000 and checked_docs <= sample_docs:
+                        ids = tokenizer.encode(value)
+                        sampled_tokens += len(ids)
+                        sampled_chars += len(value)
+                        unk = getattr(tokenizer, "unk_id", None)
+                        if unk is not None:
+                            unk_tokens += sum(1 for i in ids if i == unk)
+
+                    if remaining <= 0:
+                        break
+            else:
+                data = f.read(min(max(remaining, 0), sample_bytes))
+                remaining -= len(data)
+                checked_docs += 1
+
+                try:
+                    data.decode("utf-8")
+                except UnicodeDecodeError as exc:
+                    if exc.start < len(data) - 4:
+                        bad_encoding += 1
+
+                if not data.strip():
+                    empty_docs += 1
+                elif tokenizer is not None:
+                    sample_text = data.decode("utf-8", errors="ignore")[:200000]
+                    ids = tokenizer.encode(sample_text)
+                    sampled_tokens += len(ids)
+                    sampled_chars += len(sample_text)
+                    unk = getattr(tokenizer, "unk_id", None)
+                    if unk is not None:
+                        unk_tokens += sum(1 for i in ids if i == unk)
+
+    report.stats.update({
+        "sampled_documents": checked_docs,
+        "malformed_documents": malformed,
+        "encoding_errors": bad_encoding,
+        "empty_documents": empty_docs,
+    })
+
+    if malformed:
+        report.warnings.append(f"{malformed} malformed JSON documents in sample")
+
+    if bad_encoding:
+        report.warnings.append(f"{bad_encoding} documents with invalid UTF-8 in sample")
+
+    if tokenizer is not None and sampled_tokens:
+        report.stats["chars_per_token"] = sampled_chars / sampled_tokens
+
+    if tokenizer is not None and sampled_tokens and getattr(tokenizer, "unk_id", None) is not None:
+        rate = unk_tokens / sampled_tokens
+        report.stats["unk_rate"] = rate
+
+        if rate > 0.05:
+            report.warnings.append(f"unknown-token rate {rate:.1%} in sample; tokenizer may not fit this corpus")
+
+    return report
+
+
+def validate_shards(directory, tokenizer=None, verify_checksums=False):
+    report = ValidationReport()
+    manifest_path = os.path.join(directory, MANIFEST_NAME)
+
+    if not os.path.isfile(manifest_path):
+        report.errors.append("manifest.json not found")
+        return report
+
+    with open(manifest_path) as f:
+        manifest = json.load(f)
+
+    itemsize = {"uint16": 2, "uint32": 4}.get(manifest["dtype"])
+
+    if itemsize is None:
+        report.errors.append(f"unsupported dtype {manifest['dtype']}")
+        return report
+
+    if tokenizer is not None:
+        if not identity_matches(manifest["tokenizer_identity"], tokenizer.identity):
+            report.errors.append("tokenizer identity does not match shard manifest")
+
+        if tokenizer.vocab_size != manifest["vocab_size"]:
+            report.errors.append("tokenizer vocab size does not match shard manifest")
+
+    total = 0
+
+    for shard in manifest["shards"]:
+        path = os.path.join(directory, shard["file"])
+
+        if not os.path.isfile(path):
+            report.errors.append(f"missing shard {shard['file']}")
+            continue
+
+        size = os.path.getsize(path)
+
+        if size != shard["token_count"] * itemsize:
+            report.errors.append(f"shard {shard['file']} size does not match token_count")
+            continue
+
+        total += shard["token_count"]
+
+        if verify_checksums and file_sha256(path) != shard["checksum"]:
+            report.errors.append(f"checksum mismatch in {shard['file']}")
+
+    if total != manifest["total_tokens"]:
+        report.errors.append("manifest total_tokens does not match shard token counts")
+
+    stale = [p for p in manifest.get("source_files", []) if not os.path.isfile(p)]
+
+    if stale:
+        report.warnings.append(f"{len(stale)} source files no longer exist; corpus_version cannot be re-verified")
+    elif manifest.get("source_files") and corpus_version(manifest["source_files"]) != manifest["corpus_version"]:
+        report.warnings.append("source corpus changed since shards were built")
+
+    report.stats.update({
+        "shards": len(manifest["shards"]),
+        "total_tokens": total,
+        "checksums_verified": verify_checksums,
+    })
+
+    return report