README.md

metadata

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

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

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

@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!