Efficient_legal_model_distillation

🏛️ Gemma-2 2B Distilled — Bangladesh Legal AI Assistant

This repository provides fine-tuned LoRA adapter weights and a 4-bit quantized GGUF model for domain-specific legal reasoning and question answering over Bangladesh Law.

The model was developed through Progressive Knowledge Distillation, transferring the structured legal reasoning and answer generation capabilities of a 9-Billion parameter teacher (Gemma-2 9B) into a lightweight 2-Billion parameter student (Gemma-2 2B). It is designed to work in tandem with a Hybrid RAG engine (BM25 + intfloat/multilingual-e5-small) for cited, faithful legal assistance.

---

📌 Model Summary

• Student Architecture: google/gemma-2-2b-it (2.61B parameters)
• Teacher Model: google/gemma-2-9b-it (via Ollama & Hugging Face)
• Distillation Method:
  1. Supervised Fine-Tuning (SFT) with 3-Step Chain-of-Thought (<thought> ... </thought>) structured legal reasoning.
  2. Sequence-Level Knowledge Distillation (SeqKD) with Top-50 Sparse Teacher Logits ($\sim 95%$ probability mass) and KL Divergence at temperature $T = 4.0$.
• Target Domain: Bangladesh Jurisprudence (The Constitution of Bangladesh + Comprehensive Statutory Legislation & Acts).
• Hardware Footprint: Runs comfortably in 4-bit quantization on consumer GPUs, laptops, and edge devices ($\sim 1.63\text{ GB}$ GGUF).

---

📊 Benchmark Evaluation ($N=50$)

Evaluated against teacher reference answers across $N=50$ legal benchmark queries using automated lexical/semantic metrics and LLM judge scoring (gemma3:27b):

---

🛠️ How to Use

1. Python Inference with Hugging Face & PEFT

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

BASE_MODEL_ID = "google/gemma-2-2b-it"
ADAPTER_ID    = "nafis8766/Efficient_legal_model_distillation"

# 1. Load Tokenizer & Base Model
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# 2. Load Distilled LoRA Adapter
model = PeftModel.from_pretrained(base_model, ADAPTER_ID)

# 3. Format Prompt & Generate
prompt = "What are the key constitutional safeguards against arbitrary arrest and detention in Bangladesh?"
messages = [
    {"role": "user", "content": prompt}
]

formatted_prompt = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to("cuda")

with torch.no_grad():
    outputs = model.generate(
        formatted_prompt,
        max_new_tokens=512,
        temperature=0.3,
        top_p=0.9
    )

response = tokenizer.decode(outputs[0][formatted_prompt.shape[-1]:], skip_special_tokens=True)
print(response)

---

2. Edge & Local Inference via GGUF (llama.cpp)

Download the quantized model binary gemma-2-2b-legal-q4.gguf and run locally:

llama-cli \
  -m gemma-2-2b-legal-q4.gguf \
  -p "<start_of_turn>user\nWhat is the legal procedure for filing a writ petition under Article 102 of the Bangladesh Constitution?<end_of_turn>\n<start_of_turn>model\n" \
  -n 512 \
  --temp 0.3

---

🔬 Training & Distillation Pipeline

  Bangladesh Legal Corpus (Constitution & Acts)
                     │
                     ▼
          14.5k Synthetic Queries
                     │
                     ▼
      Gemma-2 9B Teacher (3-Step CoT)
                     │
                     ▼
    Top-50 Logits Tensor (T = 4.0, ~95% Mass)
                     │
                     ▼
  Gemma-2 2B Student (Sparse KL Divergence + SFT)
                     │
                     ▼
        Distilled LoRA Adapter Weights

• LoRA Configuration: $r=64$, $\alpha=128$, target modules [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj].
• Precision & Hardware: bfloat16 / 4-bit NF4 quantized training on NVIDIA RTX 4080 Super (16 GB VRAM).
• Loss Function: Joint cross-entropy on teacher completions + sparse KL divergence on top-50 teacher vocabulary logits.

---

🔗 Related Resources

• GitHub Codebase: https://github.com/nafis8766/Lega-LLM-model-distillation
• Progress Report & Methodology: Detailed experiment logs and analysis available in the repository docs/report_27_4.md.

---

⚖️ Citation & Disclaimer

@misc{legal_llm_distillation_2026,
  title={Efficient Legal AI for Bangladesh Law via Progressive Knowledge Distillation and Hybrid RAG},
  author={Nafis, Md.},
  year={2026},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/nafis8766/Efficient_legal_model_distillation}}
}

Disclaimer: This model is developed for academic research and assistive legal information retrieval. It does not replace professional legal consultation from a certified advocate.