---
language: en
license: apache-2.0
tags:
- text-classification
- weight-retrieval
- domain-adaptation
- medical
- legal
- code
datasets:
- qiaojin/PubMedQA
- lex_glue
- code_search_net
metrics:
- accuracy
library_name: transformers
---

# WRAG 2.0: Weight-Retrieval Augmented Generation

**WRAG 2.0** is a novel neural architecture that uses dynamic weight retrieval during the forward pass to achieve domain specialization without retraining the base model.

## 🎯 Performance

Trained on 3 domains with only 500 samples per domain (5 epochs):

| Domain | Accuracy |
|--------|----------|
| Medical (PubMedQA) | **70%** |
| Legal (LexGLUE) | **84%** |
| Code (CodeSearchNet) | **95%** |
| **Average** | **83%** |

## 🚀 Quick Start

```python
from modeling_wrag2 import WRAG2TextModel
import torch

# Load model
model = WRAG2TextModel(
    base_model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    num_shards=10,
    k=3,
    num_wr_layers=3
)

# Load trained weights
state_dict = torch.load("pytorch_model.bin")
model.load_state_dict(state_dict, strict=False)

# Move to device
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

# Inference
text = ["Question: What is the treatment for diabetes?"]
inputs = model.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256).to(device)

with torch.no_grad():
    logits, shard_scores = model(**inputs)
    prediction = logits.argmax(dim=-1)
    
print(f"Prediction: {prediction.item()}")  # 0 or 1
```

## 💡 Key Features

- ✅ **Domain Specialization**: Different weight shards activate for different domains
- ✅ **Efficient**: Only 264M trainable params (20% of total model)
- ✅ **Fast Training**: 5 epochs, ~30 minutes on L4 GPU
- ✅ **Memory Efficient**: Fits in 24GB GPU
- ✅ **Lightweight Download**: Only 264MB (base model downloaded separately)

## 🏗️ Architecture

```
Input Text
    ↓
Frozen TinyLlama-1.1B (feature extraction)
    ↓
Mean Pooling
    ↓
Weight Retrieval Layer 1 (10 shards, k=3)
    ↓ ReLU
Weight Retrieval Layer 2 (10 shards, k=3)
    ↓ ReLU
Weight Retrieval Layer 3 (10 shards, k=3)
    ↓ ReLU
Classification Head (2 classes)
    ↓
Output
```

## 📊 Model Details

- **Base Model**: TinyLlama/TinyLlama-1.1B-Chat-v1.0 (frozen)
- **Trainable Parameters**: 264M (20.2% of total)
- **Total Parameters**: 1.3B
- **Weight Shards**: 10 per layer
- **Top-k Selection**: 3 shards per forward pass
- **Number of WR Layers**: 3

## 🔧 Training Details

- **Optimizer**: AdamW
- **Learning Rate**: 1e-4
- **Batch Size**: 4
- **Epochs**: 5
- **Max Sequence Length**: 256
- **Training Data**: 500 samples per domain (medical, legal, code)
- **Training Time**: ~30 minutes on NVIDIA L4 GPU

## 📦 Installation

```bash
pip install torch transformers huggingface_hub
```

## 🎓 How It Works

WRAG 2.0 uses **dynamic weight retrieval**:

1. Input text is encoded by frozen TinyLlama
2. Each weight retrieval layer:
   - Computes similarity between input and 10 weight shards
   - Selects top-3 most relevant shards
   - Composes dynamic weight matrix via weighted combination
   - Applies dynamic weights to input
3. Classification head produces final prediction

This allows the model to specialize for different domains without retraining the base model!

## ⚠️ Limitations

- Binary classification only (yes/no, 0/1)
- Trained on limited data (500 samples per domain)
- Requires TinyLlama base model (~4.4GB) to be downloaded separately
- Best for domain-specific tasks similar to training domains

## 📝 Citation

```bibtex
@misc{wrag2-2025,
  title={WRAG 2.0: Weight-Retrieval Augmented Generation},
  year={2025},
  url={https://huggingface.co/YOUR_USERNAME/wrag2-text-classifier}
}
```

## 📄 License

Apache 2.0

## 🤝 Contributing

Issues and pull requests welcome!