Metadata-Version: 2.4
Name: qreflex
Version: 1.0.2
Summary: QreFLEX (reFLEX) — a small transformer with a gated, trainable experience-retrieval side channel: reason first, consult learned experience when useful, reason again.
Author: Qarvexium
License-Expression: Apache-2.0
Project-URL: Homepage, https://huggingface.co/qvx-o
Project-URL: Repository, https://huggingface.co/Qarvexium
Keywords: transformer,language-model,retrieval-augmented-generation,experience-replay,reflex
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Operating System :: OS Independent
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Requires-Dist: tokenizers>=0.15
Requires-Dist: tqdm>=4.60
Provides-Extra: hf
Requires-Dist: huggingface_hub>=0.19; extra == "hf"
Provides-Extra: serve
Requires-Dist: fastapi>=0.100; extra == "serve"
Requires-Dist: uvicorn>=0.23; extra == "serve"
Requires-Dist: pydantic>=2.0; extra == "serve"
Provides-Extra: all
Requires-Dist: huggingface_hub>=0.19; extra == "all"
Requires-Dist: fastapi>=0.100; extra == "all"
Requires-Dist: uvicorn>=0.23; extra == "all"
Requires-Dist: pydantic>=2.0; extra == "all"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: httpx>=0.24; extra == "dev"
Requires-Dist: build>=1.0; extra == "dev"
Dynamic: license-file

# QreFLEX

**QreFLEX** — **Q**uery-**re**trieving **F**lexible **L**earning and **EX**perience — is a small experimental conversational language model designed around a modular architecture that separates the conversational process into distinct components.

Rather than relying entirely on one large language model, QreFLEX divides responsibilities across three specialized components:

- **Main** — The primary language-generation network
- **Experience** — A lightweight retrieval/context component that surfaces potentially useful information from previous interactions
- **Intent** — A small component that helps interpret the nature of the input

The model employs a three-phase forward pass: an initial reasoning phase over the immediate context, a consultation phase where the model retrieves relevant past interactions from a non-parametric memory bank, and a final reasoning phase that integrates the retrieved experience with the current context. This design allows the model to learn and draw upon a dynamic, updatable memory without requiring full re-training of the parametric weights.

## What is QreFLEX?

QreFLEX is primarily an experimental conversational model architecture, not a factual knowledge model. It explores a different question than traditional large language models:

> How much conversational behavior can emerge from a modular architecture when some responsibilities are moved outside the main language generator?

The model is designed for:
- Experimental conversational AI
- Research into modular language-model architectures
- Studying retrieval-augmented generation
- Exploring dynamic, updatable memory systems
- Educational experimentation with transformer architectures

### Intended Use

QreFLEX is particularly interesting for conversational interactions where the model can:
- Draw upon previous interactions stored in its experience bank
- Adapt its responses based on retrieved context
- Build up a dynamic memory over time without retraining

### Limitations

QreFLEX is an experimental architecture with important limitations:
- Responses quality depends on model size and training data
- The experience bank requires periodic rebuilding and compression
- Not designed as a factual question-answering system
- Should not be used for medical, legal, financial, or safety-critical applications

The model's behavior should be understood as emerging from its training data and architecture, not as human-like understanding.

## Installation

You can install QreFLEX via pip for standard usage:

```bash
pip install qreflex
```

Installation registers four console commands: `qreflex-train`, `qreflex-eval`, `qreflex-generate`, and `qreflex-serve`.

## Architecture overview

The QreFLEX architecture introduces a retrieval-augmented component directly within the transformer forward pass, with a modular design that separates different aspects of language understanding and generation.

### Conceptual Structure

```
Input
  │
  ▼
┌─────────┐
│ Intent  │  ← Interprets input nature
└────┬────┘
  │
  ▼
┌──────────┐
│Experience│  ← Retrieves relevant context
└────┬─────┘
  │
  ▼
┌─────────┐
│  Main   │  ← Generates response
│  Model  │
└────┬────┘
  │
  ▼
Response
```

### Key Components

- **Main transformer**: Consists of causal self-attention and SwiGLU feed-forward network blocks divided into two stages. Stage 1 processes input without external experience. Stage 2 incorporates experience via cross-attention.

- **IntentEncoder**: A specialized small module that produces a query component from the input context, helping the model understand what kind of information might be relevant.

- **ExperienceEncoder**: Encodes historical interactions into (key, value) pairs to be stored in the memory bank. This is a separate network from the main model, optimized for producing good retrieval representations.

- **ExperienceBank**: A non-parametric associative memory that stores (key, value) pairs from past interactions. It can be periodically rebuilt and compressed to maintain efficiency and relevance. This is the model's dynamic, updatable memory.

- **Experience gate**: A learned scalar (initialized to zero) that controls the contribution of retrieved experience to the main network via a hyperbolic tangent (tanh) gating function. This ensures experience integration starts at zero and gradually learns when to use retrieved context.

- **Contrastive retrieval loss**: An auxiliary loss applied during training to ensure the query encoder (used by the main model) and key encoder (used by the experience encoder) learn aligned representations, making retrieval effective.

### Parameter Distribution

The QreFLEX architecture allows flexible parameter allocation. A typical small configuration (~14M total parameters) might distribute as:

| Component | Parameters | Purpose |
|-----------|-----------|---------|
| Main model | ~13M | Language generation |
| Experience encoder | ~0.8M | Encoding interactions for retrieval |
| Intent encoder | ~0.2M | Query generation |

This means the majority of capacity remains dedicated to language generation, while relatively small components provide retrieval and contextualization.

## Python API

The `QReFLEX` class provides a high-level API for model operations.

### Quick start

Training and generating with a QreFLEX model:

```python
from qreflex import QReFLEX

# Create and train a model
model = QReFLEX()
model.train(data=["Hello, world!", "Another training example."])

# Generate text
response = model.generate("Hello,")
print(response)

# Have a conversation
reply = model.chat("What's up?", history=["Hello", "Hi there!"])
print(reply)
```

### Loading a pretrained model

```python
from qreflex import QReFLEX

# Load from a local checkpoint directory
checkpoint_file = os.path.join(model_path, "reFLEX-v1-50M.pt")
tokenizer_path = os.path.join(model_path, "tokenizer.json")
model = QReFLEX.load(checkpoint_file, tokenizer=tokenizer_path)

# Or with auto-discovery of tokenizer
model = QReFLEX.load("path/to/checkpoint.pt")
```

### Loading from HuggingFace Hub

Use `from_hf()` to load models directly from HuggingFace:

```python
from qreflex import QReFLEX

# Load from HuggingFace (uses HF temp cache, not permanently saved)
model = QReFLEX.from_hf("username/model-name")

# Download ENTIRE repo and save locally
model = QReFLEX.from_hf("username/model-name", save=True)
# ✓ Full repo downloaded to: ./models/username_model-name/

# Next time with save=True, it auto-loads from local files (instant!)
model = QReFLEX.from_hf("username/model-name", save=True)
# ✓ Loading from local cache: ./models/username_model-name/

# Load specific checkpoint (works with any .pt filename)
model = QReFLEX.from_hf(
    repo_id="username/model-name",
    checkpoint_filename="my-custom-checkpoint.pt",
    save=True,
    save_dir="./my_models"
)

# Once loaded, use like any other model
response = model.generate(
    "How are you?",
    use_experience=True,
    temperature=0.7
)
```

The `from_hf()` method:
- **Smart caching**: Auto-loads from local files if previously downloaded with `save=True`
- **Full repo download**: When `save=True`, downloads the entire repository (all files)
- **Flexible filenames**: No hardcoded filename restrictions - finds any `.pt` file automatically
- **Temp cache mode**: When `save=False` (default), uses HuggingFace's temporary cache
- Returns a fully functional `QReFLEX` instance ready for `.generate()`, `.chat()`, etc.

**Note**: Loading from HuggingFace requires the `huggingface_hub` package:
```bash
pip install huggingface_hub
```

**Note**: Loading from HuggingFace requires the `huggingface_hub` package:
```bash
pip install huggingface_hub
```

### Model configuration

The `ModelConfig` class defines the model architecture. Preset sizes include `tiny`, `small`, `base`, and `large`. Custom configurations can be specified manually.

**Preset Sizes:**
```python
ModelConfig(size="tiny")    # ~1-2M parameters
ModelConfig(size="small")   # ~14M parameters  
ModelConfig(size="base")    # ~20M parameters
ModelConfig(size="large")   # ~35M parameters
```

**Custom Configuration:**
```python
ModelConfig(
    size="custom",
    dim=256,              # Main model dimension
    layers=6,             # Number of main layers
    heads=8,              # Attention heads
    ffn_dim=704,          # Feed-forward dimension
    side_dim=96,          # Experience/Intent dimension
    side_layers=2,        # Experience/Intent layers
    side_heads=4,         # Experience/Intent heads
    side_ffn_dim=256,     # Experience/Intent FFN dimension
    seq_len=256,          # Sequence length
    vocab_size=8000,      # Vocabulary size
)
```

**Key Fields:**
- `size`: Preset size name or "custom"
- `dim`: Main model dimension (must be divisible by heads)
- `layers`: Number of main transformer layers
- `heads`: Number of attention heads
- `ffn_dim`: Feed-forward network dimension
- `side_dim`: Experience/Intent encoder dimension
- `side_layers`: Number of experience/intent layers
- `side_heads`: Experience/Intent attention heads
- `seq_len`: Maximum sequence length
- `vocab_size`: Tokenizer vocabulary size

The `TrainConfig` class manages training hyperparameters.

**Common Configuration:**
```python
TrainConfig(
    steps=5000,           # Total training steps
    batch_size=16,        # Batch size
    learning_rate=3e-4,   # Peak learning rate
    warmup_steps=300,     # Warmup steps
    grad_accum=1,         # Gradient accumulation steps
    log_every=50,         # Logging interval
    save_every=1000,      # Checkpoint save interval
)
```

**Key Fields:**
- `steps`: Total training steps (not epochs)
- `batch_size`: Sequences per batch
- `learning_rate`: Peak learning rate
- `warmup_steps`: Number of warmup steps
- `weight_decay`: Weight decay coefficient
- `grad_clip`: Gradient clipping value
- `grad_accum`: Gradient accumulation steps
- `log_every`: Log metrics every N steps
- `save_every`: Save checkpoint every N steps
- `streaming`: Use streaming mode for large datasets
- `bank_max_size`: Maximum experience bank size
- `bank_rebuild_every`: Rebuild bank every N steps

### Training

The `train()` method of the `QReFLEX` class supports training from lists of strings, dictionaries, or a path to a JSONL file. 

```python
model.train(data="path/to/dataset.jsonl", config=TrainConfig(steps=5000))
```

It supports streaming mode for large datasets, validation splits, gradient accumulation, and deterministic seeding.

### Building a tokenizer

Use `build_tokenizer()` to train a new tokenizer on your dataset:

```python
model.build_tokenizer(data="path/to/dataset.jsonl", vocab_size=32000)
```

### Conversational data

Helper functions for processing dialogue datasets into the format expected by the model:

```python
from qreflex import conversations_to_texts, save_dataset

# Convert conversation turns into formatted text
conversations = [
    ["Hello", "Hi there!", "How are you?", "I'm good!"],
    ["What's up?", "Not much, you?"]
]

texts = conversations_to_texts(conversations, speakers=("A", "B"))
# Result: ["A: Hello\nB: Hi there!\nA: How are you?\nB: I'm good!", ...]

# Save to JSONL for training
save_dataset(texts, "conversations.jsonl")
```

### Conversational behavior

An interesting characteristic of QreFLEX models is that responses can vary based on the conversational context and the experience bank's contents. The experience retrieval mechanism allows the model to:

- Surface relevant previous interactions
- Maintain conversational continuity across sessions
- Adapt responses based on accumulated experience

Example behavior with experience enabled:

```python
# First interaction
model.remember("A: What's your favorite color?\nB: I love blue.")

# Later, the model can retrieve this context
response = model.generate("A: What colors do you like?\nB:", use_experience=True)
# May reference the earlier interaction about blue
```

### Experience ablation

Testing with and without the experience component shows how retrieval influences generation:

```python
# Without experience (pure language model)
response_baseline = model.generate(prompt, use_experience=False)

# With experience (retrieval-augmented)
response_with_exp = model.generate(prompt, use_experience=True)
```

The experience mechanism should be understood as a way to influence generation through contextual information rather than as a conventional knowledge database.

### Generation

The `generate()` method produces text given a prompt. The `chat()` method is tailored for conversational interactions, maintaining interaction history. Both methods accept generation parameters such as `temperature`, `top_k`, and `max_new_tokens`.

```python
# Basic generation
model.generate("Explain quantum computing:", temperature=0.7)

# Generation with stop tokens - stops immediately when these tokens are generated
model.generate(
    "List three colors:\n1.",
    max_new_tokens=50,
    stop_tokens=[".", "\n\n"]  # Stops at period or double newline
)

# Chat with custom parameters
model.chat(
    "What's the weather like?",
    history=["Hello", "Hi there!"],
    temperature=0.8,
    top_k=30
)
```

**Stop tokens**: When `stop_tokens` is provided as a list of strings, generation stops immediately when any of those tokens are generated, preventing wasted computation. The stop token itself is excluded from the returned text.

### Evaluation

The `evaluate()` method calculates metrics on a dataset. It reports perplexity (both with and without experience consultation) and retrieval alignment. It supports streaming evaluation for large validation sets. It returns a dictionary of metrics.

### Memory management

- `remember(text)`: Manually inserts an interaction into the memory bank.
- `compress_memory(max_size)`: Reduces the memory bank to a specified maximum size, retaining the most salient experiences.
- `memory_size`: Property returning the current number of items in the memory bank.

### Saving and loading

- `save(path)`: Serializes the model, tokenizer, and experience bank to a directory.
- `QReFLEX.load(path)`: Reinstantiates a model from a saved directory. It automatically discovers and loads the associated tokenizer.

### Lower-level access

For advanced usage, core components can be imported directly: `ReFLEX`, `ReFLEXConfig`, `ExperienceBank`, etc. The `load_checkpoint` function from `qreflex.train` provides granular control over restoring training states.

## Command-line interface

### qreflex-train

Trains a model from the command line.

| Flag | Description |
|---|---|
| `--data` | Path to the training dataset (JSONL). |
| `--model-size` | Model preset size (e.g., `tiny`, `small`). |
| `--output-dir` | Directory to save checkpoints and final model. |
| `--batch-size` | Batch size. |
| `--total-steps` | Total number of training steps. |

Example:
```bash
qreflex-train --data train.jsonl --model-size small --output-dir checkpoints/ --total-steps 5000
```

### qreflex-eval

Evaluates a model. Reports perplexity with and without experience, and retrieval alignment.

| Flag | Description |
|---|---|
| `--model-path` | Path to the saved model directory. |
| `--data` | Path to the evaluation dataset (JSONL). |

### qreflex-generate

Generates text from a prompt or initiates an interactive chat.

| Flag | Description |
|---|---|
| `--model-path` | Path to the saved model directory. |
| `--prompt` | Input prompt for single-shot generation. |
| `--interactive` | Launch an interactive chat session. |

In single-shot mode, the `\n` sequence in the prompt string is unescaped to a literal newline.

### qreflex-serve

Starts an HTTP API server for the model.

| Flag | Description |
|---|---|
| `--model-path` | Path to the saved model directory. |
| `--host` | Host address to bind (default: `127.0.0.1`). |
| `--port` | Port to bind (default: `8000`). |
| `--bank-save-dir`| Directory for saving the memory bank safely. |

Endpoints:
- `GET /health`: Returns server status.
- `POST /generate`: Accepts a JSON payload with a `prompt` and parameters, returns generated text.
- `POST /learn`: Submits new interactions to the experience bank.
- `POST /bank/compress`: Triggers memory bank compression.
- `POST /bank/save`: Saves the current state of the memory bank to disk.

*Note:* The `--bank-save-dir` parameter enforces security by restricting where memory bank files can be written.

## Data preparation

The `scripts/` directory includes parsers to convert common datasets into the JSONL format expected by QreFLEX:
- **Cornell Movie Dialogs**: Parses movie script conversations.
- **DailyDialog**: Parses multi-turn daily conversations.
- **OpenAssistant OASST1**: Processes assistant interaction data.
- **PersonaChat**: Handles chit-chat dataset with persona grounding.

A `prepare_tokenizer.py` script is also provided for standalone tokenizer training.

*Note: The `scripts/` directory is not included in the pip-installed package. You must clone the repository to access these tools.*

## Testing

To run the test suite, use pytest:

```bash
pytest tests/
```

The test suite covers unit tests for core model components, API functionality, memory management operations, and training routines.

## Checkpointing

During training, checkpoints are saved to the specified output directory. A checkpoint directory contains:
- The main model state dictionary.
- The experience encoder state dictionary.
- The optimizer state.
- The learning rate scheduler state.
- The current global step.
- The model configuration.
- The path to the tokenizer.
- The serialized experience bank.
- Metadata (qreflex version, git commit, timestamp).

Training can be resumed smoothly from any valid checkpoint directory.

## License

Apache 2.0

---

## Acknowledgments

QreFLEX is an experimental architecture exploring modular approaches to conversational AI. The model is intentionally designed to be small and transparent, making it suitable for research, education, and experimentation with retrieval-augmented generation techniques.
