qreflex 0.0.1__py3-none-any.whl

qreflex/__init__.py

@@ -0,0 +1,38 @@
+"""
+QreFLEX (reFLEX) — Responsive Flexible Learning and EXperience.
+
+A small transformer whose central reasoning stream (Main) can selectively
+consult two side channels, Intent (what is this input about) and
+Experience (have I seen something like this before), without either
+channel being able to override Main's own reasoning. See the project
+README for the full architecture writeup and the "what changed from the
+original prototype" history of real bugs found and fixed by actually
+running this code against real data.
+"""
+
+from qreflex.model import (
+    ReFLEX,
+    ReFLEXConfig,
+    IntentEncoder,
+    ExperienceEncoder,
+    ExperienceBank,
+    count_parameters,
+    small_reflex_config,
+)
+from qreflex.losses import retrieval_contrastive_loss
+from qreflex.api import reFLEX
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "reFLEX",
+    "ReFLEX",
+    "ReFLEXConfig",
+    "IntentEncoder",
+    "ExperienceEncoder",
+    "ExperienceBank",
+    "count_parameters",
+    "small_reflex_config",
+    "retrieval_contrastive_loss",
+    "__version__",
+]

qreflex/api.py

@@ -0,0 +1,311 @@
+"""
+High-level "load it and generate" API for reFLEX models.
+
+    from qreflex import reFLEX
+
+    model = reFLEX(directory="./my-reflex-model", from_pretrained=True)
+    answer = model.generate("A: hello there\\nB:")
+
+    # streaming
+    for chunk in model.generate("A: hello there\\nB:", stream=True):
+        print(chunk, end="", flush=True)
+
+This sits on top of the lower-level pieces that already exist in this
+package (qreflex.model, qreflex.train.load_checkpoint,
+qreflex.data.load_tokenizer) and adds two things neither the CLI
+(qreflex-generate) nor the raw classes give you directly:
+
+  1. `directory` can be a local folder OR a Hugging Face Hub repo id
+     ("YourUser/your-reflex-model") — resolved automatically.
+  2. `from_pretrained=False` gives you a randomly-initialized model
+     (still needs a tokenizer.json, since a model's vocab isn't
+     meaningful without one) for when you want to point a fresh model
+     at a directory before/instead of loading trained weights.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from typing import Generator, Optional, Union
+
+import torch
+import torch.nn.functional as F
+
+from qreflex.model import ReFLEX, ReFLEXConfig, ExperienceBank, small_reflex_config
+from qreflex.data import load_tokenizer
+
+_CHECKPOINT_NAMES = ("reFLEX-v1-15M.pt",)
+_TOKENIZER_NAMES = ("tokenizer.json",)
+_CONFIG_NAMES = ("config.json",)
+
+# "user/repo-name" or "user/repo-name:revision" — deliberately conservative;
+# anything that doesn't match this and isn't a local path is rejected with a
+# clear error rather than silently trying (and failing) a Hub download.
+_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*$")
+
+
+def _looks_like_hf_repo_id(s: str) -> bool:
+    if os.path.exists(s) or s.startswith((".", "/", "~")):
+        return False
+    return bool(_HF_REPO_ID_RE.match(s))
+
+
+def _find_file(local_dir: str, preferred_names, contains: str, exts, kind: str) -> Optional[str]:
+    """Look for one of `preferred_names` first; otherwise fall back to the
+    single file in `local_dir` whose name contains `contains` and whose
+    extension is in `exts`. Raises if that fallback is ambiguous."""
+    for name in preferred_names:
+        p = os.path.join(local_dir, name)
+        if os.path.isfile(p):
+            return p
+    candidates = sorted(
+        f for f in os.listdir(local_dir)
+        if f.lower().endswith(exts) and contains in f.lower()
+    )
+    if len(candidates) == 1:
+        return os.path.join(local_dir, candidates[0])
+    if len(candidates) > 1:
+        raise FileNotFoundError(
+            f"Found multiple possible {kind} files in '{local_dir}': {candidates}. "
+            f"Rename the one you want to one of {preferred_names}, or pass it explicitly "
+            f"via the matching *_file argument."
+        )
+    return None
+
+
+class reFLEX:
+    """Load a reFLEX checkpoint and generate text from it.
+
+    Args:
+        directory: A local folder, or a Hugging Face Hub repo id
+            (e.g. "YourUser/your-reflex-model"). Should contain a
+            tokenizer.json, and (if from_pretrained=True) a checkpoint
+            .pt file in the format written by qreflex.train.save_checkpoint.
+        from_pretrained: True loads model weights from a checkpoint found
+            in `directory`. False builds a randomly-initialized model
+            instead (a config.json in `directory` is used if present,
+            otherwise a small default config sized to the tokenizer).
+        device: "cuda" / "cpu" / etc. Defaults to cuda if available.
+        use_experience: Default for whether generate() consults the
+            experience bank stored in the checkpoint (if any). Can be
+            overridden per-call.
+        checkpoint_file / tokenizer_file / config_file: explicit paths,
+            for when auto-detection in `directory` is ambiguous or the
+            files use non-standard names.
+    """
+
+    def __init__(
+        self,
+        directory: str,
+        from_pretrained: bool = True,
+        device: Optional[str] = None,
+        use_experience: bool = True,
+        checkpoint_file: Optional[str] = None,
+        tokenizer_file: Optional[str] = None,
+        config_file: Optional[str] = None,
+    ):
+        self.directory = directory
+        self.from_pretrained = from_pretrained
+        self.use_experience = use_experience
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+        local_dir = self._resolve_directory(directory)
+
+        tok_path = tokenizer_file or _find_file(
+            local_dir, _TOKENIZER_NAMES, "tokenizer", (".json",), "tokenizer"
+        )
+        if tok_path is None:
+            raise FileNotFoundError(
+                f"No tokenizer.json found in '{local_dir}'. A reFLEX model always needs the "
+                f"tokenizer it was (or will be) trained with — point `directory` at a folder "
+                f"that has one, or pass `tokenizer_file=` explicitly."
+            )
+        self.tokenizer = load_tokenizer(tok_path)
+
+        self.bank: Optional[ExperienceBank] = None
+        self.step = 0
+        self._experience_encoder = None
+
+        if from_pretrained:
+            ckpt_path = checkpoint_file or _find_file(
+                local_dir, _CHECKPOINT_NAMES, "checkpoint", (".pt", ".bin"), "checkpoint"
+            )
+            if ckpt_path is None:
+                raise FileNotFoundError(
+                    f"from_pretrained=True but no checkpoint (.pt) file found in '{local_dir}'. "
+                    f"Pass `checkpoint_file=` explicitly, or use from_pretrained=False to start "
+                    f"from a randomly-initialized model instead."
+                )
+            # Imported lazily: qreflex.train pulls in the training loop's deps
+            # (DataLoader etc.) that a pure inference user shouldn't need to
+            # pay for just to import qreflex.
+            from qreflex.train import load_checkpoint
+
+            self.model, self._experience_encoder, self.bank, self.step = load_checkpoint(
+                ckpt_path, device=self.device
+            )
+        else:
+            cfg_path = config_file or _find_file(
+                local_dir, _CONFIG_NAMES, "config", (".json",), "config"
+            )
+            if cfg_path is not None:
+                with open(cfg_path) as f:
+                    cfg = ReFLEXConfig(**json.load(f))
+            else:
+                cfg = small_reflex_config(vocab_size=self.tokenizer.get_vocab_size())
+            self.model = ReFLEX(cfg).to(self.device)
+
+        self.model.eval()
+
+    # -- directory / hub resolution ------------------------------------------
+
+    def _resolve_directory(self, directory: str) -> str:
+        if os.path.isdir(directory):
+            return os.path.abspath(directory)
+
+        if _looks_like_hf_repo_id(directory):
+            try:
+                from huggingface_hub import snapshot_download
+            except ImportError as e:
+                raise ImportError(
+                    "huggingface_hub is required to load a reFLEX model straight from the "
+                    "Hugging Face Hub. Install it with `pip install huggingface_hub`."
+                ) from e
+            return snapshot_download(repo_id=directory, allow_patterns=["*.json", "*.pt", "*.bin"])
+
+        raise FileNotFoundError(
+            f"'{directory}' is neither an existing local directory nor a Hugging Face Hub "
+            f"repo id (e.g. 'YourUser/your-reflex-model'). Point `directory` at a local "
+            f"folder or a public/private Hub repo."
+        )
+
+    # -- generation ------------------------------------------------------------
+
+    def _encode_prompt(self, prompt: str) -> torch.Tensor:
+        ids = [self.model.cfg.bos_token_id] + self.tokenizer.encode(prompt).ids
+        if len(ids) > self.model.cfg.max_seq_len:
+            ids = ids[-self.model.cfg.max_seq_len:]  # keep the most recent context
+        return torch.tensor([ids], dtype=torch.long, device=self.device)
+
+    @staticmethod
+    def _sample_next(logits: torch.Tensor, temperature: float, top_k: Optional[int]) -> torch.Tensor:
+        if temperature <= 0:
+            return logits.argmax(dim=-1, keepdim=True)
+        logits = logits / temperature
+        if top_k is not None:
+            k = min(top_k, logits.size(-1))
+            values, _ = torch.topk(logits, k)
+            threshold = values[:, -1].unsqueeze(-1)
+            logits = logits.masked_fill(logits < threshold, float("-inf"))
+        probs = F.softmax(logits, dim=-1)
+        return torch.multinomial(probs, num_samples=1)
+
+    def generate(
+        self,
+        prompt: str,
+        max_new_tokens: int = 64,
+        temperature: float = 0.8,
+        top_k: int = 50,
+        use_experience: Optional[bool] = None,
+        stream: bool = False,
+        stop: Optional[Union[str, list]] = None,
+    ) -> Union[str, Generator[str, None, None]]:
+        """Generate a continuation for `prompt`.
+
+        stream=False (default): returns the full generated text as one string.
+        stream=True: returns a generator yielding text chunks as they're produced.
+
+        stop: one string, or a list of strings, that ends generation as soon
+        as any of them appears in the output. The stop string itself is cut
+        from the returned/streamed text (e.g. stop="\\nA:" on a model that
+        drifts into writing the next turn will give you just that turn's
+        reply, with "\\nA:" and everything after it removed).
+        """
+        exp = self.use_experience if use_experience is None else use_experience
+        if stop is None:
+            stop_strs = None
+        elif isinstance(stop, str):
+            stop_strs = [stop]
+        else:
+            stop_strs = list(stop)
+
+        gen = self._generate(prompt, max_new_tokens, temperature, top_k, exp, stop_strs)
+        if stream:
+            return gen
+        return "".join(gen)
+
+    @torch.no_grad()
+    def _generate(
+        self, prompt, max_new_tokens, temperature, top_k, use_experience, stop_strs
+    ) -> Generator[str, None, None]:
+        self.model.eval()
+        tokens = self._encode_prompt(prompt)
+        prompt_len = tokens.size(1)
+        eos_id = self.model.cfg.eos_token_id
+        emitted = ""  # text already released to the caller so far
+
+        # A stop string can straddle a token boundary (e.g. "A" comes out one
+        # step, ":" the next). If we released "A" the instant it appeared,
+        # we couldn't take it back once "A:" completes the match one step
+        # later. So once text could still be the start of a stop string, we
+        # withhold up to (longest stop string - 1) trailing characters until
+        # either the match completes (cut it, done) or a later character
+        # rules it out (release it, it was never part of a stop string).
+        hold_back = (max(len(s) for s in stop_strs) - 1) if stop_strs else 0
+
+        for _ in range(max_new_tokens):
+            context = tokens[:, -self.model.cfg.max_seq_len:]
+            output = self.model(
+                context,
+                experience_bank=self.bank if use_experience else None,
+                use_experience=use_experience,
+            )
+            next_token = self._sample_next(output["logits"][:, -1], temperature, top_k)
+            tokens = torch.cat([tokens, next_token], dim=1)
+
+            if next_token.item() == eos_id:
+                break
+
+            # Re-decode the whole generated span each step rather than each
+            # token id in isolation — some merges only resolve to the right
+            # text once neighboring tokens are known.
+            full_text = self.tokenizer.decode(tokens[0, prompt_len:].tolist())
+
+            if stop_strs:
+                hits = [i for i in (full_text.find(s) for s in stop_strs) if i != -1]
+                if hits:
+                    cut_at = min(hits)
+                    delta = full_text[len(emitted):cut_at]
+                    if delta:
+                        yield delta
+                    return  # stop string reached — it and everything after is discarded
+
+                safe_upto = max(0, len(full_text) - hold_back)
+                delta = full_text[len(emitted):safe_upto]
+                if delta:
+                    emitted = full_text[:safe_upto]
+                    yield delta
+            else:
+                delta = full_text[len(emitted):]
+                if delta:
+                    emitted = full_text
+                    yield delta
+
+        # Generation ended (eos or max_new_tokens) without ever completing a
+        # stop-string match — release whatever was withheld as a possible
+        # prefix, since we now know for certain it wasn't one.
+        if stop_strs:
+            final_text = self.tokenizer.decode(tokens[0, prompt_len:].tolist())
+            delta = final_text[len(emitted):]
+            if delta:
+                yield delta
+
+    def __repr__(self) -> str:
+        n_params = sum(p.numel() for p in self.model.parameters())
+        return (
+            f"reFLEX(directory={self.directory!r}, from_pretrained={self.from_pretrained}, "
+            f"device={self.device!r}, params={n_params:,}, step={self.step}, "
+            f"bank_size={len(self.bank) if self.bank else 0})"
+        )

qreflex/data.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from typing import Iterator, List, Optional
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from tokenizers import Tokenizer, models, pre_tokenizers, trainers, decoders
+
+
+SPECIAL_TOKENS = ["<pad>", "<bos>", "<eos>", "<unk>"]
+
+
+def train_tokenizer(
+    text_files: List[str],
+    vocab_size: int,
+    save_path: str,
+) -> Tokenizer:
+    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = decoders.ByteLevel()
+
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        special_tokens=SPECIAL_TOKENS,
+        show_progress=True,
+    )
+    tokenizer.train(text_files, trainer)
+
+    os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+    tokenizer.save(save_path)
+    return tokenizer
+
+
+def load_tokenizer(path: str) -> Tokenizer:
+    return Tokenizer.from_file(path)
+
+
+@dataclass
+class PackedExample:
+    input_ids: List[int]
+    labels: List[int]
+    experience_tokens: List[int]
+
+
+class PackedTextDataset(Dataset):
+    """
+    Tokenizes a list of raw text documents, packs them into fixed-length
+    sequences for LM training, and separately produces an `experience_tokens`
+    field per example: the SAME window of tokens (not shifted), used as the
+    input to the ExperienceEncoder for that example.
+
+    Using the example's own tokens as its experience-encoder input, paired
+    with the query built from a prefix of that same example inside the model,
+    is what makes the contrastive retrieval loss meaningful: the model learns
+    to build queries that point at keys describing the content actually
+    relevant to it, rather than being trained to point at a fixed lookup
+    that has no relation to the current input.
+    """
+
+    def __init__(
+        self,
+        texts: List[str],
+        tokenizer: Tokenizer,
+        seq_len: int,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+    ):
+        self.tokenizer = tokenizer
+        self.seq_len = seq_len
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+
+        self.examples: List[PackedExample] = []
+        self._pack(texts)
+
+    def _pack(self, texts: List[str]):
+        buffer: List[int] = []
+        for text in texts:
+            ids = self.tokenizer.encode(text).ids
+            buffer.extend([self.bos_token_id] + ids + [self.eos_token_id])
+
+            while len(buffer) >= self.seq_len + 1:
+                chunk = buffer[: self.seq_len + 1]
+                buffer = buffer[self.seq_len:]
+
+                input_ids = chunk[:-1]
+                labels = chunk[1:]
+                self.examples.append(
+                    PackedExample(
+                        input_ids=input_ids,
+                        labels=labels,
+                        experience_tokens=input_ids,
+                    )
+                )
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, idx: int):
+        ex = self.examples[idx]
+        return {
+            "input_ids": torch.tensor(ex.input_ids, dtype=torch.long),
+            "labels": torch.tensor(ex.labels, dtype=torch.long),
+            "experience_tokens": torch.tensor(ex.experience_tokens, dtype=torch.long),
+        }
+
+
+def load_jsonl_texts(path: str, text_field: str = "text") -> List[str]:
+    texts = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            obj = json.loads(line)
+            texts.append(obj[text_field])
+    return texts
+
+
+def build_dataloader(
+    dataset: Dataset,
+    batch_size: int,
+    shuffle: bool = True,
+    num_workers: int = 0,
+) -> DataLoader:
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        drop_last=True,
+        pin_memory=torch.cuda.is_available(),
+    )

qreflex/eval.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+
+import torch
+
+from qreflex.train import load_checkpoint
+from qreflex.data import load_tokenizer, load_jsonl_texts, PackedTextDataset, build_dataloader
+
+
+@torch.no_grad()
+def evaluate_perplexity(model, experience_encoder, bank, dataloader, device, use_experience=True, max_batches=None):
+    model.eval()
+    total_loss = 0.0
+    total_tokens = 0
+    n_batches = 0
+
+    for batch in dataloader:
+        if max_batches is not None and n_batches >= max_batches:
+            break
+        tokens = batch["input_ids"].to(device)
+        targets = batch["labels"].to(device)
+
+        out = model(
+            tokens,
+            targets=targets,
+            experience_bank=bank if use_experience else None,
+            use_experience=use_experience,
+        )
+        n_valid = (targets != -100).sum().item()
+        total_loss += out["lm_loss"].item() * n_valid
+        total_tokens += n_valid
+        n_batches += 1
+
+    avg_loss = total_loss / max(1, total_tokens)
+    ppl = math.exp(min(avg_loss, 20))
+    return avg_loss, ppl
+
+
+@torch.no_grad()
+def evaluate_retrieval_alignment(model, experience_encoder, dataloader, device, max_batches=20):
+    """
+    Measures whether the query built for an example is closer to its OWN
+    experience key than to other examples' keys in the batch (retrieval
+    rank / top-1 accuracy). This is a direct check on whether the
+    contrastive retrieval loss actually taught anything, independent of
+    downstream LM performance.
+
+    max_batches=None means no limit, matching evaluate_perplexity's
+    contract for the same parameter name.
+    """
+    model.eval()
+    experience_encoder.eval()
+
+    correct = 0
+    total = 0
+    mean_rank = 0.0
+
+    for i, batch in enumerate(dataloader):
+        if max_batches is not None and i >= max_batches:
+            break
+        tokens = batch["input_ids"].to(device)
+        exp_tokens = batch["experience_tokens"].to(device)
+
+        query, _ = model.build_query(tokens)
+        keys, _ = experience_encoder(exp_tokens)
+
+        sims = query @ keys.T  # [B, B]
+        ranks = sims.argsort(dim=-1, descending=True)
+        labels = torch.arange(tokens.size(0), device=device)
+
+        top1 = (ranks[:, 0] == labels).sum().item()
+        correct += top1
+        total += tokens.size(0)
+
+        for b in range(tokens.size(0)):
+            rank_of_true = (ranks[b] == b).nonzero(as_tuple=True)[0].item()
+            mean_rank += rank_of_true
+
+    top1_acc = correct / max(1, total)
+    mean_rank = mean_rank / max(1, total)
+    return top1_acc, mean_rank
+
+
+def main(args):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, experience_encoder, bank, step = load_checkpoint(args.checkpoint, device=device)
+    print(f"loaded checkpoint at step {step}, bank size {len(bank) if bank else 0}")
+
+    tokenizer = load_tokenizer(args.tokenizer_path)
+    texts = load_jsonl_texts(args.eval_data, text_field=args.text_field)
+    dataset = PackedTextDataset(texts, tokenizer, seq_len=model.cfg.max_seq_len)
+    dataloader = build_dataloader(dataset, batch_size=args.batch_size, shuffle=False)
+
+    print(f"eval examples: {len(dataset)}")
+
+    loss_no_exp, ppl_no_exp = evaluate_perplexity(model, experience_encoder, bank, dataloader, device, use_experience=False, max_batches=args.max_batches)
+    print(f"[no experience]  loss={loss_no_exp:.4f}  ppl={ppl_no_exp:.2f}")
+
+    if bank is not None and len(bank) > 0:
+        loss_exp, ppl_exp = evaluate_perplexity(model, experience_encoder, bank, dataloader, device, use_experience=True, max_batches=args.max_batches)
+        print(f"[with experience] loss={loss_exp:.4f}  ppl={ppl_exp:.2f}")
+        print(f"delta (negative = experience helps): {loss_exp - loss_no_exp:+.4f}")
+
+        top1_acc, mean_rank = evaluate_retrieval_alignment(model, experience_encoder, dataloader, device, max_batches=args.max_batches)
+        print(f"[retrieval alignment] top-1 self-match accuracy={top1_acc:.4f}  mean_rank={mean_rank:.2f}")
+    else:
+        print("no experience bank in checkpoint; skipping experience-conditioned eval")
+
+    results = {
+        "step": step,
+        "loss_no_experience": loss_no_exp,
+        "ppl_no_experience": ppl_no_exp,
+    }
+    if bank is not None and len(bank) > 0:
+        results["loss_with_experience"] = loss_exp
+        results["ppl_with_experience"] = ppl_exp
+        results["retrieval_top1_accuracy"] = top1_acc
+        results["retrieval_mean_rank"] = mean_rank
+
+    if args.output_json:
+        with open(args.output_json, "w") as f:
+            json.dump(results, f, indent=2)
+        print(f"wrote results to {args.output_json}")
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description="Evaluate a reFLEX checkpoint")
+    p.add_argument("--checkpoint", type=str, required=True)
+    p.add_argument("--tokenizer_path", type=str, required=True)
+    p.add_argument("--eval_data", type=str, required=True)
+    p.add_argument("--text_field", type=str, default="text")
+    p.add_argument("--batch_size", type=int, default=16)
+    p.add_argument("--max_batches", type=int, default=None)
+    p.add_argument("--output_json", type=str, default=None)
+    return p
+
+
+def main_cli():
+    """Entry point for the qreflex-eval console script."""
+    parser = build_arg_parser()
+    args = parser.parse_args()
+    main(args)
+
+
+if __name__ == "__main__":
+    main_cli()