QreFLEX

QreFLEX — Query-retrieving Flexible Learning and EXperience — 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:

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:

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

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:

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:

pip install huggingface_hub

Note: Loading from HuggingFace requires the huggingface_hub package:

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.

Field	Type	Description
vocab_size	int	Vocabulary size of the tokenizer.
d_model	int	Dimensionality of the model representations.
n_layers	int	Total number of transformer layers.
n_heads	int	Number of attention heads.
stage1_layers	int	Number of layers before experience integration.
d_ff	int	Dimensionality of the feed-forward network.
dropout	float	Dropout probability.

The TrainConfig class manages training hyperparameters.

Field	Type	Default	Description
batch_size	int	16	Number of sequences per batch.
learning_rate	float	1e-4	Peak learning rate.
epochs	int	3	Number of training epochs.
warmup_steps	int	1000	Number of warmup steps for the scheduler.
max_seq_len	int	512	Maximum sequence length for inputs.
grad_accum_steps	int	1	Gradient accumulation steps.
seed	int	42	Random seed for reproducibility.
val_split	float	0.1	Fraction of data to use for validation.

Training

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

model.train(data="path/to/dataset.jsonl", config=TrainConfig(epochs=5))

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:

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:

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:

# 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:

# 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.

# 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.
--epochs	Number of training epochs.

Example:

qreflex-train --data train.jsonl --model-size small --output-dir checkpoints/

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:

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 and epoch.
• The model configuration.
• The path to the tokenizer.
• The serialized experience bank.

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.