README.md

---
license: mit
language:
- en
library_name: transformers
tags:
- lora
- peft
- reinforcement-learning
- domain-adaptation
- sentence-embeddings
- curriculum-learning
- multi-task-learning
- rag
- information-retrieval
- cross-domain
- sentence-transformers
base_model:
- sentence-transformers/all-MiniLM-L6-v2
- EphAsad/FireDevourerEmbedder-RL-v3.6
pipeline_tag: sentence-similarity
datasets:
- sentence-transformers/stsb
- nyu-mll/multi_nli
- quora
- google-research-datasets/paws
- nyu-mll/glue
- GBaker/MedQA-USMLE-4-options-hf
- lex_glue
- gbharti/finance-alpaca
- scientific_papers
model-index:
- name: DomainEmbedder-v2.6
  results:
  - task:
      type: domain-classification
      name: Domain Classification
    metrics:
    - type: accuracy
      value: 0.925
      name: Training Accuracy
    - type: accuracy
      value: 0.56
      name: Stress-Test Accuracy
---

# DomainEmbedder-v2.6

> **High-Information-Density Embeddings for Cross-Domain RAG and Retrieval**

DomainEmbedder-v2.6 produces **information-dense embeddings** optimized for retrieval-augmented generation (RAG) and cross-domain similarity matching. It combines a multi-task base embedder with domain-adaptive LoRA routing.

## What This Model Does

| Component | Description |
|-----------|-------------|
| **Base Embedder** | FireDevourerEmbedder-RL-v3.6 trained on 5 NLP tasks with RL-based task weighting |
| **Domain LoRAs** | 5 specialized adapters (Medical, Legal, Code, Finance, Scientific) |
| **RL Policy** | Automatically selects the optimal domain adapter for any input |

**Why this matters for RAG/Retrieval:**
- Embeddings encode multiple facets of meaning (similarity, entailment, paraphrase, questions)
- Domain routing provides context-appropriate representations
- Results in more precise retrieval across diverse content types

## Key Innovation: Dual RL Architecture

| Stage | RL Application | Purpose |
|-------|---------------|---------|
| Base Model Training | Task Weight Policy | Dynamically balance 5 NLP objectives during training |
| Domain Extension | Adapter Selection Policy | Route to appropriate domain LoRA at inference |

This dual RL approach is novel: **RL at training time AND inference time**.

## Quick Start

### Installation

```bash
pip install torch transformers peft
```

### Loading the Model

```python
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from peft import PeftModel

# Device setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')

# Define the base embedder architecture
class FireDevourerEmbedder(nn.Module):
    def __init__(self, base_model_name='sentence-transformers/all-MiniLM-L6-v2'):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(base_model_name)
        self.hidden_size = 384

        # Task heads
        self.sts_head = nn.Sequential(nn.Linear(384, 1), nn.Sigmoid())
        self.nli_head = nn.Linear(384, 3)
        self.qqp_head = nn.Linear(384, 2)
        self.paws_head = nn.Linear(384, 2)
        self.domain_head = nn.Linear(384, 5)

    def mean_pool(self, token_embeddings, attention_mask):
        mask = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)

    def forward(self, input_ids, attention_mask, task='encode'):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        embedding = self.mean_pool(outputs.last_hidden_state, attention_mask)

        if task == 'encode':
            return embedding
        elif task == 'domain':
            return self.domain_head(embedding)
        # Add other tasks as needed

# Define RL Policy Network
class RLPolicyNetwork(nn.Module):
    def __init__(self, input_dim=384, hidden_dim=128, num_actions=5):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )
        self.policy_head = nn.Linear(hidden_dim, num_actions)
        self.value_head = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        features = self.network(x)
        policy = torch.softmax(self.policy_head(features), dim=-1)
        value = self.value_head(features)
        return policy, value

# Load model
model_dir = "path/to/DomainEmbedder-v2.6"

# 1. Load base model with checkpoint
base_model = FireDevourerEmbedder()
checkpoint = torch.load(f"{model_dir}/FireDevourerEmbedder-RL-v3.6.pt", map_location=device)
base_model.load_state_dict(checkpoint['model_state_dict'], strict=False)
base_model.to(device)
base_model.eval()

# 2. Load RL policy
rl_policy = RLPolicyNetwork()
rl_checkpoint = torch.load(f"{model_dir}/rl_policy.pt", map_location=device)
rl_policy.load_state_dict(rl_checkpoint['policy_state_dict'])
rl_policy.to(device)
rl_policy.eval()

# 3. Load LoRA adapters (example: medical)
from peft import PeftModel
lora_model = PeftModel.from_pretrained(
    base_model.encoder,
    f"{model_dir}/medical_lora"
)
```

### Computing Embeddings with Domain Selection

```python
def get_domain_embedding(text, base_model, rl_policy, lora_models, tokenizer, device):
    """Get domain-aware embedding for input text."""
    # Tokenize
    inputs = tokenizer(text, return_tensors='pt', padding=True,
                       truncation=True, max_length=512).to(device)

    # Get base embedding
    with torch.no_grad():
        base_emb = base_model(inputs['input_ids'], inputs['attention_mask'], task='encode')

        # Get domain selection from RL policy
        policy_probs, _ = rl_policy(base_emb)
        domain_idx = torch.argmax(policy_probs, dim=-1).item()

    domains = ['medical', 'legal', 'code', 'finance', 'scientific']
    selected_domain = domains[domain_idx]
    confidence = policy_probs[0, domain_idx].item()

    return {
        'embedding': base_emb,
        'domain': selected_domain,
        'confidence': confidence,
        'all_probs': policy_probs[0].cpu().numpy()
    }

# Example usage
result = get_domain_embedding(
    "What are the symptoms of diabetes?",
    base_model, rl_policy, None, tokenizer, device
)
print(f"Domain: {result['domain']} (confidence: {result['confidence']:.2%})")
```

## Architecture

```
Input Text
    │
    ▼
┌────────────────────────────────────────────┐
│  MiniLM-L6-v2 Encoder (FROZEN)             │
│  + Optional LoRA Adapter (domain-specific) │
│  384-dimensional output                     │
└────────────────────────────────────────────┘
    │
    ├──────────────────────────────────────────┐
    │                                          │
    ▼                                          ▼
┌─────────────────┐                  ┌──────────────────┐
│  Base Embedding │                  │   RL Policy Net  │
│    (384-dim)    │                  │   (66K params)   │
└─────────────────┘                  └──────────────────┘
                                              │
                                              ▼
                                     Domain Selection
                                     [Medical, Legal, Code,
                                      Finance, Scientific]
                                              │
                                              ▼
                              Load corresponding LoRA adapter
                                              │
                                              ▼
                              Domain-Adapted Embedding
```

### Component Details

| Component | Specification |
|-----------|---------------|
| Base Encoder | MiniLM-L6-v2 (22M params) |
| Embedding Dim | 384 |
| LoRA Rank | 16 |
| LoRA Alpha | 32 |
| LoRA Target | Query, Value projections |
| LoRA Params | 147,456 per adapter (0.645%) |
| RL Policy | 66,566 params |
| Domains | Medical, Legal, Code, Finance, Scientific |

## Performance

### Base Model: Multi-Task Embedding Quality

The base FireDevourerEmbedder achieves **0.71 average** across 5 distinct NLP tasks:

| Task | Dataset | Score | What It Measures |
|------|---------|-------|------------------|
| Question Similarity | QQP | 0.8636 | Intent matching |
| Paraphrase Detection | PAWS | 0.8459 | Adversarial robustness |
| Paraphrase Detection | MRPC | 0.7744 | News domain paraphrase |
| NLI | MultiNLI | 0.7465 | Logical relationships |
| Semantic Similarity | STS-B | 0.3366 | Fine-grained similarity |
| **Average** | | **0.7134** | **Cross-task capability** |

**Philosophy**: Individual task scores are traded for cross-domain information density. This makes embeddings more versatile for RAG and retrieval across diverse content.

### Domain Routing Accuracy

**Training Results (In-Distribution)**

| Metric | Value |
|--------|-------|
| Domain Accuracy | 92.5% |
| Average Reward | 1.527 |
| Training Steps | 5,000 |

**Stress-Test Benchmark (Semantically Similar Cross-Domain Phrases)**

The benchmark intentionally uses complex, semantically similar phrases across domains to test robustness:

| Metric | DomainEmbedder (RL+LoRA) | Base Model | Improvement |
|--------|--------------------------|------------|-------------|
| Domain Accuracy | 56.0% | 20.4% | **+35.6%** |
| Avg Confidence | 28.5% | 77.6% | More calibrated |

### Per-Domain Breakdown

| Domain | DomainEmbedder | Base Model | Note |
|--------|----------------|------------|------|
| Finance | 78.0% | 0.0% | +78.0% |
| Medical | 73.0% | 0.0% | +73.0% |
| Legal | 53.0% | 15.0% | +38.0% |
| Scientific | 48.0% | 1.0% | +47.0% |
| Code | 28.0% | 86.0% | Base over-predicted code |

**Key Insight**: The base model had an 86% "code" prediction bias with high confidence. The RL+LoRA system corrects this by providing balanced, calibrated domain distribution.

## Training Details

### Domain Training Data

| Domain | Samples | Sources |
|--------|---------|---------|
| Medical | 40,000 | MedQA-USMLE, MedQuAD, PubMedQA, Medical Meadow, ChatDoctor |
| Legal | 40,000 | EUR-LEX, CaseHold, ECTHR-A, ECTHR-B |
| Code | 40,000 | Code Alpaca, MBPP, Code Contests, Python Instructions |
| Finance | 40,000 | Finance Alpaca, FinGPT-FiQA, Financial QA |
| Scientific | 40,000 | arXiv, PubMed (87.3% real + 12.7% augmented) |
| **Total** | **200,000** | |

### LoRA Training Configuration

| Parameter | Value |
|-----------|-------|
| Epochs | 3 per domain |
| Batch Size | 32 |
| Learning Rate | 2e-4 |
| Loss | Contrastive (InfoNCE-style) |
| Trainable Params | 147,456 (0.645% of base) |
| Warmup Steps | 500 |
| Max Gradient Norm | 1.0 |

### RL Training (Supervised A2C)

| Parameter | Value |
|-----------|-------|
| Algorithm | Actor-Critic (A2C) |
| Total Steps | 5,000 |
| Episodes per Step | 5 |
| Gamma (discount) | 0.99 |
| Entropy Coef | 0.1 (high exploration) |
| Value Coef | 0.5 |
| Correctness Bonus | +1.0 |
| Correctness Penalty | -0.5 |
| Baseline Decay | 0.99 |

### Curriculum Learning Phases

| Phase | Steps | Data | Accuracy |
|-------|-------|------|----------|
| 1 (Easy) | 0-1,500 | Clear domain examples (10K) | 68.8% → 87.5% |
| 2 (Moderate) | 1,500-3,500 | Easy + ambiguous (20K) | 87.5% → 89.3% |
| 3 (Hard) | 3,500-5,000 | All data incl. hybrid (28K) | 89.3% → 92.5% |

### Training Progress

| Version | Step | Accuracy | Reward |
|---------|------|----------|--------|
| v2.1 | 500 | 68.8% | 1.100 |
| v2.2 | 1,000 | 80.1% | 1.336 |
| v2.3 | 1,500 | 87.5% | 1.454 |
| v2.4 | 2,000 | 88.9% | 1.480 |
| v2.5 | 3,000 | 89.3% | 1.507 |
| **v2.6** | **4,000** | **92.5%** | **1.527** |

## Package Contents

```
DomainEmbedder-v2.6/
├── FireDevourerEmbedder-RL-v3.6.pt   # Base model checkpoint (86.7 MB)
├── rl_policy.pt                       # Trained RL policy (0.27 MB)
├── metadata.json                      # Training metadata
├── README.md                          # This file
├── medical_lora/                      # Medical domain adapter (0.6 MB)
│   ├── adapter_config.json
│   └── adapter_model.safetensors
├── legal_lora/                        # Legal domain adapter (0.6 MB)
├── code_lora/                         # Code domain adapter (0.6 MB)
├── finance_lora/                      # Finance domain adapter (0.6 MB)
└── scientific_lora/                   # Scientific domain adapter (0.6 MB)
```

**Total Size**: ~90 MB (self-contained)

## Intended Use

### Best Use Cases

- **RAG Systems**: Domain-aware retrieval for multi-domain knowledge bases
- **Cross-Domain Search**: Finding similar content across Medical, Legal, Code, Finance, Scientific domains
- **Document Classification**: Automatic domain routing for document processing pipelines
- **Semantic Similarity**: Information-dense embeddings for precise matching
- **Multi-Domain Chatbots**: Context-appropriate responses based on detected domain

### Limitations

- **English Only**: Trained exclusively on English data
- **Max Length**: 512 tokens maximum input length
- **Domain Coverage**: 5 domains only (Medical, Legal, Code, Finance, Scientific)
- **Stress-Test Accuracy**: 56% on semantically similar cross-domain queries
- **STS-B Trade-off**: Lower fine-grained similarity (0.34) for broader task coverage

## Citation

```bibtex
@misc{domainembedder2025,
  author = {Asad, Zain},
  title = {DomainEmbedder: Domain-Adaptive Embeddings with Dual RL and LoRA},
  year = {2025},
  publisher = {Hugging Face},
  note = {Multi-task base embedder with RL-based task weighting + domain-specific LoRA adapters with curriculum learning}
}
```

## Author

**Zain Asad**

## License

MIT License