README.md

---
base_model: Qwen/Qwen2.5-7B-Instruct
library_name: peft
pipeline_tag: text-generation
license_name: proprietary
license_link: LICENSE
tags:
- base_model:adapter:Qwen/Qwen2.5-7B-Instruct
- lora
- transformers
- compliance
- nist
- control-extraction
- regulatory
---

# NIST Control Extraction LoRA Adapter

A fine-tuned LoRA adapter for **Qwen2.5-7B-Instruct** specifically designed for accurate extraction of security controls from NIST framework documents. This model eliminates hallucination issues present in the base model, ensuring precise identification of controls without mistaking control enhancements or related text as valid controls.

## Key Features

- **Accurate Control Extraction**: Precisely identifies control IDs, titles, and descriptions from framework documents
- **Reduced Hallucination**: Trained to distinguish between actual controls and control enhancements/related content
- **Fast Inference**: Processes 492-page NIST documents in ~15 minutes (vs. 27 minutes with base model)
- **Structured Output**: Returns controls in clean JSON format with `<END>` token for reliable parsing

---

## Model Details

### Model Description

This LoRA adapter enhances the Qwen2.5-7B-Instruct model for the specialized task of extracting security controls from compliance framework documents. The adapter was trained using a custom weighted loss function that penalizes false positives more heavily than false negatives, addressing the critical requirement in compliance auditing where incorrectly identifying a control is more problematic than missing one.

| Property | Value |
|----------|-------|
| **Developed by** | Rishit Sharma |
| **Model Type** | LoRA Adapter (PEFT) |
| **Base Model** | [Qwen/Qwen2.5-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct) |
| **Language** | English |
| **Domain** | Compliance & Regulatory Frameworks |
| **License** | Proprietary - No use allowed without prior permission |

---

## Model Architecture

### LoRA Configuration

| Parameter | Value |
|-----------|-------|
| **Rank (r)** | 16 |
| **Alpha** | 32 |
| **Dropout** | 0.05 |
| **Target Modules** | `q_proj`, `k_proj`, `v_proj`, `o_proj` |
| **Bias** | None |
| **Task Type** | CAUSAL_LM |

### Quantization (Training)

| Parameter | Value |
|-----------|-------|
| **Quantization** | 4-bit (QLoRA) |
| **Quant Type** | NF4 |
| **Double Quantization** | Enabled |
| **Compute Dtype** | bfloat16 |

### Special Tokens

- **`<END>`**: Custom stop token appended to outputs for reliable generation termination

---

## Intended Use

### Primary Use Case

Building autonomous compliance auditing agents that can:
- Parse and analyze framework documents (PDF/text)
- Extract structured control information automatically
- Verify deployment status of controls within an organization

### Target Users

- Compliance Officers & Auditors
- GRC (Governance, Risk, Compliance) Teams
- Security Analysts
- Organizations undergoing NIST compliance assessments

---

## Quick Start

### Installation

```bash
pip install transformers peft torch accelerate bitsandbytes
```

### Loading the Model

```python
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import torch

# Quantization config (optional, for memory efficiency)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("path/to/final_adapter")

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "path/to/final_adapter")
```

### Inference Example

```python
system_prompt = """You are a senior Compliance Auditor and Regulatory Analyst specialized in ISO, NIST, and statutory frameworks."""

messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": f"Analyze this text:\n\n{page_text}"}
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt"
).to(model.device)

outputs = model.generate(
    input_ids,
    max_new_tokens=512,
    temperature=0.1,
    do_sample=False
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
```

### Expected Output Format

```json
[
    {
        "control_id": "AC-1",
        "control_title": "Access Control Policy and Procedures",
        "control_desc": "Description of the control..."
    }
]
<END>
```

---

## Training Details

### Training Configuration

| Parameter | Value |
|-----------|-------|
| **Hardware** | NVIDIA RTX 5070 Ti |
| **Training Time** | ~1 hour |
| **Epochs** | 14 (with early stopping) |
| **Batch Size** | 1 (effective: 8 with gradient accumulation) |
| **Learning Rate** | 1e-5 |
| **Optimizer** | Paged AdamW 8-bit |
| **LR Scheduler** | Cosine |
| **Warmup Ratio** | 0.05 |
| **Max Gradient Norm** | 0.3 |
| **Precision** | FP16 |

### Training Data

- **Source**: NIST SP 800-53 Framework (492 pages)
- **Dataset Creation**: Custom pipeline using Gemini Pro for initial extraction, followed by manual verification
- **Data Balance**: ~60% positive samples (pages with controls), ~40% negative samples (pages without controls)
- **Format**: JSONL with chat template structure

### Custom Loss Function

A **Weighted Loss Trainer** was implemented to address the asymmetric cost of errors in compliance:

- **Positive Weight**: 2.0x (penalizes missing actual controls less than falsely identifying controls)
- **Rationale**: In compliance auditing, falsely identifying a control (hallucination) is more problematic than missing one, as it can lead to incorrect compliance assessments

```python
# Samples with controls are weighted 2x during loss computation
weights = torch.where(has_control, 2.0, 1.0)
weighted_loss = (sample_loss * weights).mean()
```

---

## 📈 Evaluation & Performance

| Metric | Base Qwen2.5-7B | This Adapter |
|--------|-----------------|--------------|
| **Processing Time (492 pages)** | ~27 minutes | ~15 minutes |
| **Hallucination Rate** | High | Minimal |
| **Control Enhancement Confusion** | Frequent | Resolved |

---

## ⚠️ Limitations & Risks

### Current Limitations

- **Framework Specificity**: Optimized primarily for NIST SP 800-53; performance on other frameworks (ISO 27001, SOC 2, etc.) may vary
- **Language**: Trained on English documents only
- **Document Format**: Best performance on well-structured PDF documents

### Known Risks

- May require additional fine-tuning for non-NIST frameworks
- Performance depends on input text quality and preprocessing
- Should be validated by human auditors for critical compliance decisions

### Future Improvements

- [ ] Training on ISO 27001/27002 frameworks
- [ ] Multi-framework support (SOC 2, HIPAA, PCI-DSS)
- [ ] Improved handling of complex document layouts
- [ ] Longer training with expanded dataset

---

## 📄 License

**Proprietary License** - This model is not available for public use without explicit prior permission from the developer.

For licensing inquiries, please contact via the channels below.

---

## 📬 Contact

| Channel | Link |
|---------|------|
| **Email** | [rishitshar36@gmail.com](mailto:rishitshar36@gmail.com) |
| **GitHub** | [github.com/rishit836](https://github.com/rishit836) |
| **Project Repository** | [control-extraction-using-llm-finetuned](https://github.com/rishit836/control-extraction-using-llm-finetuned) |

---

## 🙏 Acknowledgments

- [Qwen Team](https://github.com/QwenLM/Qwen2.5) for the excellent base model
- [Hugging Face](https://huggingface.co/) for the Transformers and PEFT libraries
- NIST for the publicly available SP 800-53 framework documentation

---

## 📚 Citation

If you use this model in your research or project, please cite:

```bibtex
@misc{sharma2026nist-control-extraction,
    title={NIST Control Extraction LoRA Adapter for Qwen2.5-7B},
    author={Sharma, Rishit},
    year={2026},
    publisher={GitHub},
    howpublished={\url{https://github.com/rishit836/control-extraction-using-llm-finetuned}}
}
```