Metadata-Version: 2.1
Name: seren-llm
Version: 0.1.2
Summary: Seren: A Self-Refining Decoder-Only LLM with GQA, RoPE, Flash Attention, and a novel two-pass self-refinement mechanism.
Home-page: https://github.com/your-org/seren
Author: Seren Authors
License: Apache-2.0
Keywords: transformer,language-model,self-refinement,grouped-query-attention,rope,flash-attention,causal-lm,nlp,deep-learning,pytorch,huggingface
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Operating System :: OS Independent
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0.0
Requires-Dist: transformers>=4.36.0
Requires-Dist: numpy>=1.22.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: isort; extra == "dev"
Requires-Dist: mypy; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx; extra == "docs"
Requires-Dist: sphinx-rtd-theme; extra == "docs"
Requires-Dist: myst-parser; extra == "docs"
Provides-Extra: train
Requires-Dist: accelerate>=0.26.0; extra == "train"
Requires-Dist: datasets>=2.14.0; extra == "train"
Requires-Dist: tokenizers>=0.15.0; extra == "train"

# Seren 🌊

**A Self-Refining Decoder-Only LLM Architecture**

Seren is a production-grade, HuggingFace-compatible transformer library built
around a novel *two-pass self-refinement* mechanism: the model generates a
first response normally, then conditions a second, improved response on a
learned projection of its own internal states.

---

## Architecture at a glance

| Feature | Detail |
|---|---|
| Attention | Grouped Query Attention (GQA) + Flash Attention 2 via `torch.nn.functional.scaled_dot_product_attention` |
| Position | Rotary Position Embedding (RoPE) |
| Normalisation | Pre-RMSNorm |
| FFN | SwiGLU (gated MLP) |
| KV Cache | Fully compatible — no tokens discarded between passes |
| Gradient Checkpointing | ✅ |
| Weight Tying | Optional (`tie_word_embeddings`) |
| Device Parallelism | 2-GPU via HuggingFace `device_map` |
| **Self-Refinement** | **Novel two-pass mechanism (see below)** |

---

## Installation

```bash
pip install seren
# or from source:
pip install -e ".[dev]"
```

**Requirements:** Python ≥ 3.8, PyTorch ≥ 2.0, Transformers ≥ 4.36

---

## Quick start

```python
from seren import SerenConfig, SerenForCausalLM

config = SerenConfig(
    vocab_size=32_000,
    hidden_size=2_048,
    num_hidden_layers=16,
    num_attention_heads=16,
    num_key_value_heads=4,    # GQA: 4 KV heads shared across 16 query heads
    intermediate_size=5_504,
    max_position_embeddings=4_096,
    self_hidden_layer=2,    # depth of refinement MLP (2 linear layers)
    alpha=1.0,                # weight for loss₁ (Pass-1 quality)
    beta=1.0,                 # weight for loss₂ (refinement contrastive)
)

model = SerenForCausalLM(config)
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
```

---

## The Self-Refinement Mechanism

Seren generates every response in **two passes** over the same sequence:

```
[System Prompt] + [User] + [Response₁] + [Response₂]
```

### Pass 1 — Baseline generation
The model runs a full forward pass on the entire sequence with no special
conditioning.  The final-layer hidden states at the **RESP1** positions are
captured.

### Refinement projection
A depth-configurable MLP (`self_hidden_layer` linear layers, each `hidden_size → hidden_size` with SiLU between them) followed by
RMSNorm processes the Pass-1 hidden states and compresses them via mean-pooling
into a single `[B, 1, hidden_size]` conditioning vector.

### Pass 2 — Refined generation
The conditioning vector is **added** to the token embeddings of the **RESP2**
positions before the second forward pass.  The model then predicts a refined
second response conditioned on both the original context and its own internal
Pass-1 representation.

### Training objective

```python
# loss₁ — forces the first response to be good
loss1 = cross_entropy(pass1_logits[resp1_positions], resp1_labels)

# loss₂ — forces the second response to improve over the first
#   CE(pass1 at resp2) = baseline CE without refinement
#   CE(pass2 at resp2) = CE with refinement → should be lower → loss₂ > 0
loss2 = cross_entropy(pass1_logits[resp2_positions], resp2_labels) \
      - cross_entropy(pass2_logits[resp2_positions], resp2_labels)

total_loss = alpha * loss1 + beta * loss2
```

---

## Training usage

```python
import torch
from seren import SerenForCausalLM, SerenConfig

model = SerenForCausalLM(SerenConfig()).cuda()

# Prepare batch: [SYS+USR+RESP1] and [RESP2]
pass1_ids    = torch.randint(0, 32000, (2, 64)).cuda()   # [B, T1]
pass2_ids    = torch.randint(0, 32000, (2, 32)).cuda()   # [B, T2]

# Labels: -100 for positions we don't want loss on
resp1_labels = pass1_ids.clone()
resp1_labels[:, :10] = -100   # mask system/user turns

resp2_labels = pass2_ids.clone()

outputs = model(
    input_ids=pass1_ids,
    response2_input_ids=pass2_ids,
    labels=resp1_labels,
    response2_labels=resp2_labels,
)

outputs.loss.backward()
print(f"loss={outputs.loss:.4f}  loss1={outputs.loss1:.4f}  loss2={outputs.loss2:.4f}")
```

---

## Inference: two-pass generation

```python
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("...")
model = SerenForCausalLM.from_pretrained("...", device_map="balanced")

context = tokenizer("User: What is 2+2?\nAssistant:", return_tensors="pt").input_ids

# ── Pass 1: generate first response ───────────────────────────────────
out1 = model.generate(
    input_ids=context,
    max_new_tokens=128,
    return_dict_in_generate=True,
    output_hidden_states=True,
)

# Extract last-layer hidden states from the last generated step
pass1_hidden = out1.hidden_states[-1][-1]   # [B, 1, H]

# ── Compute refinement conditioning vector ─────────────────────────────
refinement_ctx = model.compute_refinement_context(pass1_hidden)

# ── Pass 2: generate refined response ─────────────────────────────────
with model.with_refinement(refinement_ctx):
    out2 = model.generate(
        input_ids=out1.sequences,
        max_new_tokens=128,
    )

print(tokenizer.decode(out2[0]))
```

---

## 2-GPU model parallelism

```python
model = SerenForCausalLM.from_pretrained(
    "path/to/checkpoint",
    device_map="balanced",   # splits layers evenly across available GPUs
    torch_dtype=torch.bfloat16,
)
```

---

## Component imports

```python
from seren import SerenConfig           # configuration
from seren import SerenForCausalLM     # full model with LM head + self-refinement
from seren import SerenModel           # transformer body only
from seren import SerenDecoderLayer    # single decoder layer
from seren import SerenRMSNorm         # RMS normalisation
from seren import SerenRotaryEmbedding # RoPE
from seren import SerenCausalLMOutputWithPast  # typed output dataclass
```

---

## License

Apache 2.0
