pytensorforge 0.1.0__py3-none-any.whl

src/core/Vector.py

@@ -0,0 +1,31 @@
+class Vector:
+    def __init__(self, data):
+        self.data = data
+
+    def set(self, index, value):
+        self.data[index] = value
+
+    def get(self, index):
+        return self.data[index]
+
+    @staticmethod
+    def load(self, data):
+        self.data = data
+
+    def dump(self):
+        return self.data
+
+    @staticmethod
+    def Zero(self, data):
+        vec = Vector(data)
+        for i in vec.data:
+            vec.set(i, 0)
+        return vec
+
+    def zero(self):
+        for i in range(len(self.data)):
+            self.data[i] = 0
+
+    def log(self):
+        for i in range(len(self.data)):
+            print(self.data[i])

src/data/chat_dataset.py

@@ -0,0 +1,188 @@
+import json
+from collections import deque
+
+import numpy as np
+
+PACKING_MODES = ("pack", "pad")
+
+
+class ChatRecordError(ValueError):
+    pass
+
+
+class ChatSFTDataset:
+
+    def __init__(self, corpus, template, sequence_length, packing="pack", messages_field="messages", state=None):
+        if packing not in PACKING_MODES:
+            raise ValueError(f"unknown packing '{packing}'; choose one of {PACKING_MODES}")
+
+        bad = [f for f in corpus.files if not f.lower().endswith(".jsonl")]
+
+        if bad:
+            raise ValueError(f"chat datasets must be .jsonl files; got {bad[0]}")
+
+        self.corpus = corpus
+        self.template = template
+        self.tokenizer = template.tokenizer
+        self.sequence_length = int(sequence_length)
+        self.packing = packing
+        self.messages_field = messages_field
+        self.roles = set(template.template.roles)
+        self.pad_id = self.tokenizer.eos_id if self.tokenizer.eos_id is not None else 0
+        self.load_state_dict(state or {})
+
+    def _records(self):
+        for file_index in range(self.file_index, len(self.corpus.files)):
+            path = self.corpus.files[file_index]
+            offset = self.byte_offset if file_index == self.file_index else 0
+
+            with open(path, "rb") as f:
+                f.seek(offset)
+
+                while True:
+                    raw = f.readline()
+
+                    if not raw:
+                        break
+
+                    end = f.tell()
+                    line = raw.strip()
+
+                    if line:
+                        yield file_index, end, self._parse(line, path, end)
+                    else:
+                        yield file_index, end, None
+
+    def _parse(self, line, path, end):
+        where = f"{path} (record ending at byte {end})"
+
+        try:
+            record = json.loads(line)
+        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+            raise ChatRecordError(f"{where}: invalid JSON: {exc}") from None
+
+        messages = record.get(self.messages_field) if isinstance(record, dict) else None
+
+        if not isinstance(messages, list) or not messages:
+            raise ChatRecordError(f"{where}: expected a non-empty '{self.messages_field}' list")
+
+        for i, m in enumerate(messages):
+            if not isinstance(m, dict) or not isinstance(m.get("role"), str) or not isinstance(m.get("content"), str):
+                raise ChatRecordError(f"{where}: message {i} needs string 'role' and 'content'")
+
+            if m["role"] not in self.roles:
+                raise ChatRecordError(
+                    f"{where}: message {i} has role '{m['role']}', which the chat template does not define"
+                )
+
+        return [{"role": m["role"], "content": m["content"]} for m in messages]
+
+    def _conversations(self):
+        for file_index, end, messages in self._records():
+            self.file_index = file_index
+            self.byte_offset = end
+
+            if messages is None:
+                continue
+
+            ids, learn = self.template.training_tokens(messages)
+
+            if not any(learn):
+                self.skipped += 1
+                continue
+
+            self.records += 1
+            yield ids, learn
+
+    def _pack(self, conversations, needed):
+        while True:
+            while len(self.tokens) < needed:
+                item = next(conversations, None)
+
+                if item is None:
+                    return None
+
+                ids, learn = item
+                self.tokens.extend(ids)
+                self.learn.extend(learn)
+
+            block = [self.tokens.popleft() for _ in range(needed)]
+            mask = [self.learn.popleft() for _ in range(needed)]
+            return block, mask
+
+    def _pad(self, conversations, needed):
+        while True:
+            item = next(conversations, None)
+
+            if item is None:
+                return None
+
+            ids, learn = item
+            ids, learn = ids[:needed], learn[:needed]
+
+            if len(ids) < 2 or not any(learn[1:]):
+                self.truncated_away += 1
+                continue
+
+            fill = needed - len(ids)
+            return ids + [self.pad_id] * fill, learn + [0] * fill
+
+    def batches(self, batch_size):
+        needed = self.sequence_length + 1
+        conversations = self._conversations()
+        take = self._pack if self.packing == "pack" else self._pad
+
+        try:
+            while True:
+                rows = []
+
+                for _ in range(batch_size):
+                    row = take(conversations, needed)
+
+                    if row is None:
+                        break
+
+                    rows.append(row)
+
+                if not rows:
+                    return
+
+                ids = np.array([r[0] for r in rows], dtype=np.int64)
+                learn = np.array([r[1] for r in rows], dtype=np.float32)
+                yield ids[:, :-1], ids[:, 1:], learn[:, 1:]
+        finally:
+            conversations.close()
+
+    def start_new_epoch(self):
+        self.load_state_dict({})
+
+    def state_dict(self):
+        return {
+            "kind": "chat",
+            "packing": self.packing,
+            "file_index": self.file_index,
+            "byte_offset": self.byte_offset,
+            "token_buffer": list(self.tokens),
+            "learn_buffer": list(self.learn),
+            "records": self.records,
+            "skipped": self.skipped,
+            "truncated_away": self.truncated_away,
+        }
+
+    def load_state_dict(self, data):
+        if data and data.get("kind") != "chat":
+            raise ValueError("dataset state is not from a chat dataset")
+
+        if data and data.get("packing", self.packing) != self.packing:
+            raise ValueError("dataset state was written with a different packing mode")
+
+        self.file_index = int(data.get("file_index", 0))
+        self.byte_offset = int(data.get("byte_offset", 0))
+        self.tokens = deque(data.get("token_buffer", []))
+        self.learn = deque(data.get("learn_buffer", []))
+        self.records = int(data.get("records", 0))
+        self.skipped = int(data.get("skipped", 0))
+        self.truncated_away = int(data.get("truncated_away", 0))
+
+        if len(self.tokens) != len(self.learn):
+            raise ValueError("corrupt chat dataset state: token and mask buffers differ in length")

src/data/corpus.py

@@ -0,0 +1,104 @@
+import os
+import random
+
+
+class CorpusIndex:
+
+    def __init__(
+        self,
+        paths,
+        extensions=(".txt", ".jsonl", ".json"),
+        recursive=True,
+        shuffle=False,
+        seed=0,
+    ):
+        self.extensions = tuple(e.lower() for e in extensions)
+        self.recursive = recursive
+        self.shuffle = shuffle
+        self.seed = seed
+
+        self.files = self._discover(paths)
+
+        if shuffle:
+            rng = random.Random(seed)
+            rng.shuffle(self.files)
+
+    def _discover(self, paths):
+        if isinstance(paths, str):
+            paths = [paths]
+
+        found = []
+
+        for path in paths:
+            path = os.path.abspath(path)
+
+            if os.path.isfile(path):
+                if path.lower().endswith(self.extensions):
+                    found.append(path)
+                continue
+
+            if not os.path.isdir(path):
+                raise FileNotFoundError(f"corpus path not found: {path}")
+
+            if self.recursive:
+                for root, _, names in os.walk(path):
+                    for name in sorted(names):
+                        if name.lower().endswith(self.extensions):
+                            found.append(os.path.join(root, name))
+            else:
+                for name in sorted(os.listdir(path)):
+                    full = os.path.join(path, name)
+                    if os.path.isfile(full) and name.lower().endswith(self.extensions):
+                        found.append(full)
+
+        found.sort()
+
+        return found
+
+    def __len__(self):
+        return len(self.files)
+
+    def state_dict(self):
+        return {
+            "files": list(self.files),
+            "shuffle": self.shuffle,
+            "seed": self.seed,
+        }
+
+    def load_state_dict(self, state):
+        if state["files"] != self.files:
+            raise ValueError(
+                "checkpoint corpus file list does not match the current "
+                "corpus; resuming would silently skip or duplicate data"
+            )
+
+    def split(self, validation_fraction, seed=0):
+        if not 0 < validation_fraction < 1:
+            raise ValueError("validation_fraction must be between 0 and 1")
+
+        import hashlib
+
+        train, val = [], []
+
+        for path in self.files:
+            h = hashlib.sha256(f"{seed}:{os.path.basename(path)}".encode()).digest()
+            bucket = int.from_bytes(h[:8], "big") / 2**64
+            (val if bucket < validation_fraction else train).append(path)
+
+        if not val and len(train) > 1:
+            val.append(train.pop())
+
+        if not train:
+            raise ValueError("split left no training files; use a smaller validation_fraction or more files")
+
+        a = CorpusIndex.__new__(CorpusIndex)
+        b = CorpusIndex.__new__(CorpusIndex)
+
+        for obj, files in ((a, train), (b, val)):
+            obj.extensions = self.extensions
+            obj.recursive = self.recursive
+            obj.shuffle = self.shuffle
+            obj.seed = self.seed
+            obj.files = files
+
+        return a, b

src/data/document_stream.py

@@ -0,0 +1,178 @@
+import json
+import os
+
+JSON_EXTENSIONS = (".jsonl", ".json")
+ASCII_WS = frozenset(b" \t\n\r\x0b\x0c")
+WS_BYTES = (b" ", b"\t", b"\n", b"\r", b"\x0b", b"\x0c")
+
+
+def _complete_utf8_len(buf):
+    n = len(buf)
+    i = n - 1
+    back = 0
+
+    while i >= 0 and back < 4 and (buf[i] & 0xC0) == 0x80:
+        i -= 1
+        back += 1
+
+    if i < 0:
+        return n
+
+    lead = buf[i]
+
+    if lead < 0x80:
+        need = 1
+    elif lead >> 5 == 0b110:
+        need = 2
+    elif lead >> 4 == 0b1110:
+        need = 3
+    elif lead >> 3 == 0b11110:
+        need = 4
+    else:
+        need = 1
+
+    return n if n - i >= need else i
+
+
+def _prev_char_is_space(buf, end):
+    k = end - 1
+
+    while k > 0 and (buf[k] & 0xC0) == 0x80:
+        k -= 1
+
+    try:
+        return buf[k:end].decode("utf-8").isspace()
+    except UnicodeDecodeError:
+        return False
+
+
+def _find_cut(buf):
+    limit = len(buf)
+
+    while True:
+        last = max(buf.rfind(c, 0, limit) for c in WS_BYTES)
+
+        if last <= 0:
+            return -1
+
+        start = last
+
+        while start > 0 and buf[start - 1] in ASCII_WS:
+            start -= 1
+
+        if start == 0:
+            return -1
+
+        if not _prev_char_is_space(buf, start):
+            return start
+
+        limit = start
+
+
+class DocumentReader:
+
+    def __init__(self, files, read_buffer_size=1 << 20, text_field="text"):
+        self.files = files
+        self.read_buffer_size = read_buffer_size
+        self.text_field = text_field
+
+    def _read_txt(self, f, start_offset):
+        buf = b""
+        offset = start_offset
+
+        while True:
+            chunk = f.read(self.read_buffer_size)
+
+            if not chunk:
+                if buf:
+                    offset += len(buf)
+                    yield buf.decode("utf-8", errors="ignore"), offset, True
+                return
+
+            buf += chunk
+
+            cut = _find_cut(buf)
+
+            if cut <= 0:
+                if len(buf) < 4 * self.read_buffer_size:
+                    continue
+
+                cut = _complete_utf8_len(buf)
+
+                if cut >= len(buf):
+                    cut = _complete_utf8_len(buf[:-1])
+
+                if cut <= 0:
+                    continue
+
+            segment = buf[:cut]
+            buf = buf[cut:]
+            offset += len(segment)
+
+            text = segment.decode("utf-8", errors="ignore")
+
+            if text:
+                yield text, offset, False
+
+    def _read_jsonl(self, f, start_offset):
+        buf = b""
+        offset = start_offset
+
+        while True:
+            chunk = f.read(self.read_buffer_size)
+
+            if not chunk:
+                if buf.strip():
+                    text = self._extract_json(buf)
+                    offset += len(buf)
+                    if text:
+                        yield text, offset, True
+                break
+
+            buf += chunk
+
+            while b"\n" in buf:
+                line, buf = buf.split(b"\n", 1)
+                offset += len(line) + 1
+
+                if line.strip():
+                    text = self._extract_json(line)
+                    if text:
+                        yield text, offset, True
+
+    def _extract_json(self, raw_line):
+        try:
+            obj = json.loads(raw_line.decode("utf-8", errors="ignore"))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            return None
+
+        if isinstance(obj, dict):
+            value = obj.get(self.text_field, "")
+            return value if isinstance(value, str) else None
+
+        return None
+
+    def read_file(self, path, start_offset=0):
+        ext = os.path.splitext(path)[1].lower()
+
+        with open(path, "rb") as f:
+            if start_offset:
+                f.seek(start_offset)
+
+            if ext in JSON_EXTENSIONS:
+                yield from self._read_jsonl(f, start_offset)
+            else:
+                yield from self._read_txt(f, start_offset)
+
+    def iter_documents(self, start_file_index=0, start_byte_offset=0):
+        file_index = start_file_index
+        resume_offset = start_byte_offset
+
+        while file_index < len(self.files):
+            path = self.files[file_index]
+            offset = resume_offset if file_index == start_file_index else 0
+
+            for text, next_offset, doc_end in self.read_file(path, offset):
+                yield file_index, next_offset, text, doc_end
+
+            file_index += 1

src/data/parallel_encode.py

@@ -0,0 +1,86 @@
+import multiprocessing as mp
+from collections import deque
+
+import numpy as np
+
+_EMPTY = np.zeros(0, dtype=np.int32)
+_worker_tokenizer = None
+
+
+def _init_worker(tokenizer):
+    global _worker_tokenizer
+    _worker_tokenizer = tokenizer
+
+
+def _encode_texts(texts):
+    tok = _worker_tokenizer
+    return [np.asarray(tok.encode(t), dtype=np.int32) if t else _EMPTY for t in texts]
+
+
+def encode_serial(tokenizer, doc_iter):
+    for file_index, offset, text, doc_end in doc_iter:
+        ids = tokenizer.encode(text) if text else ()
+        yield file_index, offset, ids, doc_end
+
+
+class ParallelEncoder:
+
+    def __init__(self, tokenizer, workers, batch_chars=1 << 18, max_batch_docs=4096, inflight_per_worker=2,
+                 start_method="spawn"):
+        if workers < 2:
+            raise ValueError("ParallelEncoder needs at least 2 workers; use encode_serial otherwise")
+
+        if tokenizer.vocab_size > np.iinfo(np.int32).max:
+            raise ValueError("vocabulary too large for int32 token transport")
+
+        self.tokenizer = tokenizer
+        self.workers = int(workers)
+        self.batch_chars = max(1, int(batch_chars))
+        self.max_batch_docs = max(1, int(max_batch_docs))
+        self.max_inflight = max(1, self.workers * int(inflight_per_worker))
+        self.start_method = start_method
+
+    def _next_batch(self, it):
+        meta = []
+        texts = []
+        chars = 0
+
+        for file_index, offset, text, doc_end in it:
+            meta.append((file_index, offset, doc_end))
+            texts.append(text)
+            chars += len(text)
+
+            if chars >= self.batch_chars or len(texts) >= self.max_batch_docs:
+                break
+
+        return meta, texts
+
+    def encode(self, doc_iter):
+        ctx = mp.get_context(self.start_method)
+        pool = ctx.Pool(self.workers, initializer=_init_worker, initargs=(self.tokenizer,))
+        pending = deque()
+        it = iter(doc_iter)
+        exhausted = False
+
+        try:
+            while True:
+                while not exhausted and len(pending) < self.max_inflight:
+                    meta, texts = self._next_batch(it)
+
+                    if not meta:
+                        exhausted = True
+                        break
+
+                    pending.append((meta, pool.apply_async(_encode_texts, (texts,))))
+
+                if not pending:
+                    return
+
+                meta, result = pending.popleft()
+
+                for (file_index, offset, doc_end), ids in zip(meta, result.get()):
+                    yield file_index, offset, ids, doc_end
+        finally:
+            pending.clear()
+            pool.terminate()
+            pool.join()

src/data/prefetch.py

@@ -0,0 +1,62 @@
+import queue
+import threading
+
+_STOP = object()
+
+
+class PrefetchLoader:
+
+    def __init__(self, iterable_factory, prefetch_size=4):
+        self.iterable_factory = iterable_factory
+        self.prefetch_size = max(1, prefetch_size)
+        self._queue = None
+        self._thread = None
+        self._error = None
+        self._halt = threading.Event()
+
+    def _put(self, item):
+        while not self._halt.is_set():
+            try:
+                self._queue.put(item, timeout=0.1)
+                return True
+            except queue.Full:
+                continue
+
+        return False
+
+    def _worker(self):
+        try:
+            for item in self.iterable_factory():
+                if not self._put(item):
+                    return
+        except Exception as exc:
+            self._error = exc
+        finally:
+            self._put(_STOP)
+
+    def __iter__(self):
+        self._queue = queue.Queue(maxsize=self.prefetch_size)
+        self._error = None
+        self._halt.clear()
+
+        self._thread = threading.Thread(target=self._worker, daemon=True)
+        self._thread.start()
+
+        try:
+            while True:
+                item = self._queue.get()
+
+                if item is _STOP:
+                    if self._error is not None:
+                        raise self._error
+                    return
+
+                yield item
+        finally:
+            self.stop()
+
+    def stop(self):
+        self._halt.set()
+
+        if self._thread is not None and self._thread.is_alive() and self._thread is not threading.current_thread():
+            self._thread.join(timeout=1.0)