control-extraction-qwen-model

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.

---

Model Architecture

LoRA Configuration

Quantization (Training)

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

pip install transformers peft torch accelerate bitsandbytes

Loading the Model

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

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

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

---

Training Details

Training Configuration

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

# 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

---

⚠️ 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

---

🙏 Acknowledgments

• Qwen Team for the excellent base model
• Hugging Face 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:

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