Clean lightweight deployment

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,38 @@
+# Use an official Python runtime as a parent image
+FROM python:3.10-slim
+
+# Set the working directory to /app
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    git \
+    git-lfs \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy the requirements file into the container
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Pre-download HF Models to cache them in the Docker image (prevents downloading on every boot)
+RUN python -c "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM; AutoTokenizer.from_pretrained('nsi319/legal-pegasus'); AutoModelForSeq2SeqLM.from_pretrained('nsi319/legal-pegasus');"
+RUN python -c "from sentence_transformers import SentenceTransformer, CrossEncoder; SentenceTransformer('BAAI/bge-base-en-v1.5'); CrossEncoder('BAAI/bge-reranker-base');"
+
+# Copy the rest of the application code (including .git if not ignored)
+COPY . .
+
+
+# Download heavy databases from Dataset into their correct folders
+RUN huggingface-cli download SaiPranav09/NyayLens-Data paragraphs.db --repo-type dataset --local-dir data/processed/indexed/
+RUN huggingface-cli download SaiPranav09/NyayLens-Data faiss_index.bin --repo-type dataset --local-dir data/processed/faiss/
+RUN huggingface-cli download SaiPranav09/NyayLens-Data model.safetensors --repo-type dataset --local-dir outputs/summarization/final/
+
+
+# Expose port 7860 (Hugging Face standard)
+EXPOSE 7860
+
+# Command to run the application using Uvicorn
+CMD ["uvicorn", "src.api.main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,11 @@
+---
+title: NyayLens API
+emoji: 📚
+colorFrom: purple
+colorTo: purple
+sdk: docker
+pinned: false
+short_description: NyayLens-Legal AI Assistant
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

outputs/summarization/final/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f3beb2c13e1b5773787c8bb39605511469c4f77d96d2f38cca086c14d352f37
+size 437956140

outputs/summarization/final/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbda47393ea48d83251e18b1d9e470b83cb2fa30ca734d2095ff10fc6b4fd2c6
+size 5368

requirements.txt

@@ -0,0 +1,27 @@
+# NyayLens — Python Backend Dependencies
+# Install with: pip install -r requirements.txt
+
+# API
+fastapi==0.115.12
+uvicorn==0.34.0
+python-multipart==0.0.20
+pydantic==2.12.5
+
+# RAG & Retrieval
+faiss-cpu==1.9.0
+sentence-transformers==5.2.0
+groq==1.0.0
+
+# Summarization
+transformers==4.57.3
+torch==2.6.0
+sentencepiece==0.2.1
+huggingface-hub==0.36.0
+
+# PDF Extraction
+pdfplumber==0.11.8
+
+# Utilities
+python-dotenv==1.0.1
+safetensors==0.4.5
+ragas==0.2.14

src/api/main.py

@@ -0,0 +1,223 @@
+# src/api/main.py
+import sys
+import os
+import io
+import time
+import uuid
+import atexit
+import shutil
+import asyncio
+from pathlib import Path
+from typing import Optional
+
+from fastapi import FastAPI, HTTPException, UploadFile, File, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, field_validator
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
+
+from src.rag.query_engine import QueryEngine
+from src.summarization.inference import summarize
+
+# ── Constants ──────────────────────────────────────────────────────────────
+MAX_UPLOAD_MB   = 10
+MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024
+UPLOAD_DIR      = Path("data/uploads")
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+SUMMARIZE_TIMEOUT_S = 180   # 3 min max for summarization on CPU
+
+# ── App ────────────────────────────────────────────────────────────────────
+app = FastAPI(
+    title="NyayLens API",
+    description="Production API for Legal Chat, Document QA, and Summarization",
+    version="1.0.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[
+        "https://nyay-lens.vercel.app",  # Production Vercel URL
+        "http://localhost:5173",         # Local Vite dev server
+        "http://127.0.0.1:5173"
+    ],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ── Startup / Shutdown ─────────────────────────────────────────────────────
+async def cleanup_loop():
+    """Background task to remove leftover files older than 2 hours."""
+    while True:
+        now = time.time()
+        for f in UPLOAD_DIR.glob("*"):
+            if f.is_file() and (now - f.stat().st_mtime) > 7200:
+                try:
+                    f.unlink()
+                except Exception as e:
+                    print(f"Cleanup error: {e}")
+        await asyncio.sleep(3600)  # Check every hour
+
+@app.on_event("startup")
+async def startup():
+    global query_engine
+    print("Initializing NyayLens Backend...")
+    query_engine = QueryEngine()
+    
+    # Start the infinite cleanup loop
+    asyncio.create_task(cleanup_loop())
+    print("✓ Backend ready. Background cleanup active.")
+
+@app.on_event("shutdown")
+def shutdown():
+    """Clean up all uploaded files on server shutdown."""
+    if UPLOAD_DIR.exists():
+        shutil.rmtree(UPLOAD_DIR)
+        UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+        print("✓ Uploads directory cleaned on shutdown.")
+
+# ── Schema ─────────────────────────────────────────────────────────────────
+class UnifiedRequest(BaseModel):
+    message: str
+    filepath: Optional[str] = None
+    top_k: int = 5
+    chat_history: Optional[list] = []
+
+    @field_validator("message")
+    @classmethod
+    def message_not_empty(cls, v):
+        if not v or not v.strip():
+            raise ValueError("Message cannot be empty")
+        if len(v) > 4000:
+            raise ValueError("Message too long (max 4000 characters)")
+        return v.strip()
+
+# ── Health ─────────────────────────────────────────────────────────────────
+@app.get("/")
+@app.get("/api/health")
+def health():
+    return {
+        "status":  "online",
+        "service": "NyayLens API",
+        "version": "1.0.0",
+        "models":  ["Legal-BERT", "Legal-PEGASUS", "Llama-3.1-8B (Groq)"],
+        "index":   "FAISS 298K vectors",
+    }
+
+# ── Upload ─────────────────────────────────────────────────────────────────
+@app.post("/api/upload")
+async def upload_document(file: UploadFile = File(...)):
+    """
+    Accepts .pdf and .txt files up to 10 MB.
+    PDFs are extracted to plain text via pdfplumber.
+    Returns a server filepath for subsequent /api/chat calls.
+    """
+    import pdfplumber
+
+    # 1. Validate extension
+    filename = file.filename or "upload"
+    ext = Path(filename).suffix.lower()
+    if ext not in {".pdf", ".txt"}:
+        raise HTTPException(status_code=400, detail="Only .pdf and .txt files are supported.")
+
+    # 2. Read with size guard
+    raw_bytes = await file.read()
+    if len(raw_bytes) > MAX_UPLOAD_BYTES:
+        raise HTTPException(
+            status_code=413,
+            detail=f"File too large. Maximum allowed size is {MAX_UPLOAD_MB} MB."
+        )
+    if len(raw_bytes) == 0:
+        raise HTTPException(status_code=400, detail="Uploaded file is empty.")
+
+    # 3. Unique name to avoid collisions
+    uid       = uuid.uuid4().hex[:8]
+    safe_name = f"{uid}_{Path(filename).stem}"
+
+    # 4. Extract / save
+    if ext == ".pdf":
+        text_parts = []
+        try:
+            with pdfplumber.open(io.BytesIO(raw_bytes)) as pdf:
+                for page in pdf.pages:
+                    t = page.extract_text()
+                    if t:
+                        text_parts.append(t.strip())
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"PDF extraction failed: {e}")
+
+        if not text_parts:
+            raise HTTPException(
+                status_code=422,
+                detail="PDF contains no readable text. It may be a scanned image — please use a searchable PDF."
+            )
+
+        out_path = UPLOAD_DIR / f"{safe_name}.txt"
+        out_path.write_text("\n\n".join(text_parts), encoding="utf-8")
+        return {"filepath": str(out_path), "filename": filename, "pages": len(text_parts), "size_kb": round(len(raw_bytes)/1024, 1)}
+
+    else:
+        out_path = UPLOAD_DIR / f"{safe_name}.txt"
+        out_path.write_bytes(raw_bytes)
+        return {"filepath": str(out_path), "filename": filename, "size_kb": round(len(raw_bytes)/1024, 1)}
+
+
+# ── Chat ───────────────────────────────────────────────────────────────────
+@app.post("/api/chat")
+def chat(request: UnifiedRequest):
+    """
+    Unified intent-aware chat endpoint.
+    Routes to: Summarization | Document QA | Global RAG
+    """
+    message_lower = request.message.lower()
+
+    print(f"\n[BACKEND] '{request.message[:80]}' | file={os.path.basename(request.filepath) if request.filepath else 'None'}")
+
+    # Validate filepath if provided
+    if request.filepath:
+        if not os.path.exists(request.filepath):
+            return JSONResponse(
+                status_code=404,
+                content={"answer": "The uploaded document could not be found on the server. Please re-upload the file.", "sources": []}
+            )
+
+    try:
+        # ── Route 1: Summarization (with timeout) ──────────────────────────
+        if "summarize" in message_lower or "summary" in message_lower:
+            if not request.filepath:
+                return {
+                    "answer": "Please **upload a PDF or text file** first using the 📎 button, then ask me to summarize it.",
+                    "sources": []
+                }
+            print("[BACKEND] → Summarization pipeline")
+            summary_dict = summarize(request.filepath)
+            return {
+                "answer": "__STRUCTURED_SUMMARY__",
+                "summary": summary_dict,
+                "sources": [{"judgment_id": os.path.basename(request.filepath), "score": 1.0}]
+            }
+
+        # ── Route 2: Document QA ────────────────────────────────────────────
+        if request.filepath:
+            print("[BACKEND] → Document QA")
+            return query_engine.query_with_document(request.message, request.filepath, chat_history=request.chat_history)
+
+        # ── Route 3: Global RAG ─────────────────────────────────────────────
+        print("[BACKEND] → Global RAG")
+        return query_engine.query(request.message, top_k=request.top_k, chat_history=request.chat_history)
+
+    except Exception as e:
+        print(f"[BACKEND ERROR] {e}")
+        raise HTTPException(status_code=500, detail=f"An internal error occurred: {str(e)}")
+
+
+# ── Cleanup old uploads (files older than 2 hours) ─────────────────────────
+@app.delete("/api/upload/{filename}")
+def delete_upload(filename: str):
+    """Explicit delete for a specific upload."""
+    target = UPLOAD_DIR / filename
+    if target.exists() and target.is_file():
+        target.unlink()
+        return {"status": "deleted"}
+    raise HTTPException(status_code=404, detail="File not found.")

src/extraction/batch_processor.py

@@ -0,0 +1,459 @@
+# """
+# Batch processor for large-scale PDF extraction
+# Includes progress tracking, error handling, and resumability
+# """
+
+# from pathlib import Path
+# from typing import List, Dict
+# import logging
+# from tqdm import tqdm
+# import json
+# from datetime import datetime
+# from pdf_extractor import LegalJudgmentExtractor
+
+# logging.basicConfig(
+#     level=logging.INFO,
+#     format='%(asctime)s - %(levelname)s - %(message)s'
+# )
+# logger = logging.getLogger(__name__)
+
+
+# class BatchProcessor:
+#     """
+#     Batch processor for extracting text from thousands of legal judgments
+#     Features: Progress tracking, resumability, comprehensive reporting
+#     """
+    
+#     def __init__(self, 
+#                  input_dir: Path, 
+#                  output_dir: Path,
+#                  num_workers: int = 1,
+#                  enable_ocr: bool = True):
+        
+#         self.input_dir = Path(input_dir)
+#         self.output_dir = Path(output_dir)
+#         self.num_workers = num_workers
+        
+#         # Create extractor
+#         self.extractor = LegalJudgmentExtractor(output_dir, enable_ocr=enable_ocr)
+        
+#         # Progress tracking
+#         self.progress_file = output_dir / "processing_progress.json"
+#         self.stats_file = output_dir / "processing_stats.json"
+        
+#     def get_all_pdfs(self) -> List[Path]:
+#         """Get all PDF files from input directory"""
+#         return sorted(self.input_dir.rglob("*.pdf"))
+    
+#     def load_progress(self) -> set:
+#         """Load already processed files"""
+#         if self.progress_file.exists():
+#             with open(self.progress_file, 'r') as f:
+#                 data = json.load(f)
+#                 return set(data.get('processed_files', []))
+#         return set()
+    
+#     def save_progress(self, processed_files: set):
+#         """Save processing progress"""
+#         with open(self.progress_file, 'w') as f:
+#             json.dump({
+#                 'processed_files': list(processed_files),
+#                 'last_updated': datetime.now().isoformat(),
+#                 'total_processed': len(processed_files)
+#             }, f, indent=2)
+    
+#     def process_single_pdf(self, pdf_path: Path) -> Dict:
+#         """Process a single PDF"""
+#         try:
+#             success = self.extractor.process_pdf(pdf_path)
+#             return {
+#                 'filename': pdf_path.name,
+#                 'year': pdf_path.parent.name,
+#                 'success': success,
+#                 'error': None
+#             }
+#         except Exception as e:
+#             return {
+#                 'filename': pdf_path.name,
+#                 'year': pdf_path.parent.name,
+#                 'success': False,
+#                 'error': str(e)
+#             }
+    
+#     def process_batch(self, 
+#                      start_year: int = None, 
+#                      end_year: int = None,
+#                      limit: int = None,
+#                      resume: bool = True):
+#         """
+#         Process PDFs in batch with progress tracking
+        
+#         Args:
+#             start_year: Start from this year (inclusive)
+#             end_year: Process until this year (inclusive)
+#             limit: Maximum number of PDFs to process
+#             resume: Continue from last checkpoint
+#         """
+        
+#         logger.info("Starting batch processing...")
+#         logger.info(f"Workers: {self.num_workers}")
+        
+#         # Get all PDFs
+#         all_pdfs = self.get_all_pdfs()
+#         logger.info(f"Found {len(all_pdfs):,} PDFs")
+        
+#         # Filter by year if specified
+#         if start_year or end_year:
+#             all_pdfs = [
+#                 p for p in all_pdfs 
+#                 if (not start_year or int(p.parent.name) >= start_year) and
+#                    (not end_year or int(p.parent.name) <= end_year)
+#             ]
+#             logger.info(f"Filtered to {len(all_pdfs):,} PDFs (years {start_year}-{end_year})")
+        
+#         # Load progress and filter already processed
+#         if resume:
+#             processed = self.load_progress()
+#             all_pdfs = [p for p in all_pdfs if str(p) not in processed]
+#             logger.info(f"Resuming: {len(all_pdfs):,} PDFs remaining")
+#         else:
+#             processed = set()
+        
+#         # Apply limit
+#         if limit:
+#             all_pdfs = all_pdfs[:limit]
+#             logger.info(f"Limited to {len(all_pdfs):,} PDFs")
+        
+#         if not all_pdfs:
+#             logger.info("No PDFs to process!")
+#             return
+        
+#         # Initialize stats
+#         stats = {
+#             'total': len(all_pdfs),
+#             'successful': 0,
+#             'failed': 0,
+#             'start_time': datetime.now().isoformat(),
+#             'failed_files': []
+#         }
+        
+#         # Process with progress bar
+#         with tqdm(total=len(all_pdfs), desc="Processing PDFs") as pbar:
+#             for pdf_path in all_pdfs:
+#                 result = self.process_single_pdf(pdf_path)
+                
+#                 if result['success']:
+#                     stats['successful'] += 1
+#                 else:
+#                     stats['failed'] += 1
+#                     stats['failed_files'].append({
+#                         'file': result['filename'],
+#                         'error': result['error']
+#                     })
+#                     logger.warning(f"Failed: {result['filename']} - {result['error']}")
+                
+#                 # Update progress
+#                 processed.add(str(pdf_path))
+                
+#                 # Save progress every 50 files
+#                 if len(processed) % 50 == 0:
+#                     self.save_progress(processed)
+                
+#                 pbar.update(1)
+#                 pbar.set_postfix({
+#                     'Success': stats['successful'],
+#                     'Failed': stats['failed']
+#                 })
+        
+#         # Final save
+#         self.save_progress(processed)
+        
+#         # Save statistics
+#         stats['end_time'] = datetime.now().isoformat()
+#         stats['success_rate'] = (stats['successful'] / stats['total'] * 100) if stats['total'] > 0 else 0
+        
+#         with open(self.stats_file, 'w') as f:
+#             json.dump(stats, f, indent=2)
+        
+#         # Summary
+#         logger.info("\n" + "="*60)
+#         logger.info("PROCESSING COMPLETE")
+#         logger.info("="*60)
+#         logger.info(f"Total processed: {stats['total']:,}")
+#         logger.info(f"Successful: {stats['successful']:,}")
+#         logger.info(f"Failed: {stats['failed']:,}")
+#         logger.info(f"Success rate: {stats['success_rate']:.2f}%")
+#         logger.info("="*60)
+
+
+# def main():
+#     """Main execution with CLI arguments"""
+#     import argparse
+    
+#     parser = argparse.ArgumentParser(description='Batch process legal judgment PDFs')
+#     parser.add_argument('--start-year', type=int, help='Start year (inclusive)')
+#     parser.add_argument('--end-year', type=int, help='End year (inclusive)')
+#     parser.add_argument('--limit', type=int, help='Maximum PDFs to process')
+#     parser.add_argument('--no-resume', action='store_true', help='Start fresh')
+#     parser.add_argument('--no-ocr', action='store_true', help='Disable OCR fallback')
+    
+#     args = parser.parse_args()
+    
+#     # Configuration
+#     INPUT_DIR = Path("data/raw")
+#     OUTPUT_DIR = Path("data/processed/extracted")
+    
+#     processor = BatchProcessor(
+#         input_dir=INPUT_DIR,
+#         output_dir=OUTPUT_DIR,
+#         enable_ocr=not args.no_ocr
+#     )
+    
+#     processor.process_batch(
+#         start_year=args.start_year,
+#         end_year=args.end_year,
+#         limit=args.limit,
+#         resume=not args.no_resume
+#     )
+
+
+# if __name__ == "__main__":
+#     main()
+"""
+Batch processor for large-scale PDF extraction
+Includes progress tracking, error handling, and resumability
+"""
+
+from pathlib import Path
+from typing import List, Dict
+import logging
+from tqdm import tqdm
+import json
+from datetime import datetime
+from pdf_extractor import LegalJudgmentExtractor
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+class BatchProcessor:
+    """
+    Batch processor for extracting text from thousands of legal judgments
+    Features: Progress tracking, resumability, comprehensive reporting
+    """
+    
+    def __init__(self, 
+                 input_dir: Path, 
+                 output_dir: Path,
+                 num_workers: int = 1,
+                 enable_ocr: bool = True):
+        
+        self.input_dir = Path(input_dir)
+        self.output_dir = Path(output_dir)
+        self.num_workers = num_workers
+        
+        # Create extractor
+        self.extractor = LegalJudgmentExtractor(output_dir, enable_ocr=enable_ocr)
+        
+        # Progress tracking
+        self.progress_file = output_dir / "processing_progress.json"
+        self.stats_file = output_dir / "processing_stats.json"
+        
+    def get_all_pdfs(self) -> List[Path]:
+        """Get all PDF files (case-insensitive)"""
+        # FIX: Search for both .pdf and .PDF
+        pdfs_lower = list(self.input_dir.rglob("*.pdf"))
+        pdfs_upper = list(self.input_dir.rglob("*.PDF"))
+        all_pdfs = pdfs_lower + pdfs_upper
+        return sorted(set(all_pdfs))  # Remove duplicates and sort
+    
+    def load_progress(self) -> set:
+        """Load already processed files"""
+        if self.progress_file.exists():
+            with open(self.progress_file, 'r') as f:
+                data = json.load(f)
+                return set(data.get('processed_files', []))
+        return set()
+    
+    def save_progress(self, processed_files: set):
+        """Save processing progress"""
+        with open(self.progress_file, 'w') as f:
+            json.dump({
+                'processed_files': list(processed_files),
+                'last_updated': datetime.now().isoformat(),
+                'total_processed': len(processed_files)
+            }, f, indent=2)
+    
+    def process_single_pdf(self, pdf_path: Path) -> Dict:
+        """Process a single PDF"""
+        try:
+            success = self.extractor.process_pdf(pdf_path)
+            return {
+                'filename': pdf_path.name,
+                'year': pdf_path.parent.name,
+                'success': success,
+                'error': None
+            }
+        except Exception as e:
+            return {
+                'filename': pdf_path.name,
+                'year': pdf_path.parent.name,
+                'success': False,
+                'error': str(e)
+            }
+    
+    def process_batch(self, 
+                     start_year: int = None, 
+                     end_year: int = None,
+                     limit: int = None,
+                     resume: bool = True):
+        """
+        Process PDFs in batch with progress tracking
+        
+        Args:
+            start_year: Start from this year (inclusive)
+            end_year: Process until this year (inclusive)
+            limit: Maximum number of PDFs to process
+            resume: Continue from last checkpoint
+        """
+        
+        logger.info("Starting batch processing...")
+        logger.info(f"Input directory: {self.input_dir}")
+        logger.info(f"Output directory: {self.output_dir}")
+        
+        # Get all PDFs
+        all_pdfs = self.get_all_pdfs()
+        logger.info(f"Found {len(all_pdfs):,} PDFs")
+        
+        if len(all_pdfs) == 0:
+            logger.error("❌ No PDFs found! Check your data/raw directory.")
+            logger.error(f"Looking in: {self.input_dir}")
+            logger.error("Make sure PDFs are in year folders like: data/raw/1950/*.PDF")
+            return
+        
+        # Filter by year if specified
+        if start_year or end_year:
+            filtered_pdfs = []
+            for p in all_pdfs:
+                try:
+                    year = int(p.parent.name)
+                    if (not start_year or year >= start_year) and (not end_year or year <= end_year):
+                        filtered_pdfs.append(p)
+                except ValueError:
+                    logger.warning(f"Skipping non-year folder: {p.parent.name}")
+            
+            all_pdfs = filtered_pdfs
+            logger.info(f"Filtered to {len(all_pdfs):,} PDFs (years {start_year}-{end_year})")
+        
+        # Load progress and filter already processed
+        if resume:
+            processed = self.load_progress()
+            all_pdfs = [p for p in all_pdfs if str(p) not in processed]
+            logger.info(f"Resuming: {len(all_pdfs):,} PDFs remaining")
+        else:
+            processed = set()
+        
+        # Apply limit
+        if limit:
+            all_pdfs = all_pdfs[:limit]
+            logger.info(f"Limited to {len(all_pdfs):,} PDFs")
+        
+        if not all_pdfs:
+            logger.info("No PDFs to process!")
+            return
+        
+        # Initialize stats
+        stats = {
+            'total': len(all_pdfs),
+            'successful': 0,
+            'failed': 0,
+            'start_time': datetime.now().isoformat(),
+            'failed_files': []
+        }
+        
+        # Process with progress bar
+        with tqdm(total=len(all_pdfs), desc="Processing PDFs") as pbar:
+            for pdf_path in all_pdfs:
+                result = self.process_single_pdf(pdf_path)
+                
+                if result['success']:
+                    stats['successful'] += 1
+                else:
+                    stats['failed'] += 1
+                    stats['failed_files'].append({
+                        'file': result['filename'],
+                        'year': result['year'],
+                        'error': result['error']
+                    })
+                    logger.warning(f"Failed: {result['filename']} - {result['error']}")
+                
+                # Update progress
+                processed.add(str(pdf_path))
+                
+                # Save progress every 50 files
+                if len(processed) % 50 == 0:
+                    self.save_progress(processed)
+                
+                pbar.update(1)
+                pbar.set_postfix({
+                    'Success': stats['successful'],
+                    'Failed': stats['failed']
+                })
+        
+        # Final save
+        self.save_progress(processed)
+        
+        # Save statistics
+        stats['end_time'] = datetime.now().isoformat()
+        stats['success_rate'] = (stats['successful'] / stats['total'] * 100) if stats['total'] > 0 else 0
+        
+        with open(self.stats_file, 'w') as f:
+            json.dump(stats, f, indent=2)
+        
+        # Summary
+        logger.info("\n" + "="*60)
+        logger.info("PROCESSING COMPLETE")
+        logger.info("="*60)
+        logger.info(f"Total processed: {stats['total']:,}")
+        logger.info(f"Successful: {stats['successful']:,}")
+        logger.info(f"Failed: {stats['failed']:,}")
+        logger.info(f"Success rate: {stats['success_rate']:.2f}%")
+        logger.info("="*60)
+
+
+def main():
+    """Main execution with CLI arguments"""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Batch process legal judgment PDFs')
+    parser.add_argument('--start-year', type=int, help='Start year (inclusive)')
+    parser.add_argument('--end-year', type=int, help='End year (inclusive)')
+    parser.add_argument('--limit', type=int, help='Maximum PDFs to process')
+    parser.add_argument('--no-resume', action='store_true', help='Start fresh')
+    parser.add_argument('--no-ocr', action='store_true', help='Disable OCR fallback')
+    
+    args = parser.parse_args()
+    
+    # Configuration
+    INPUT_DIR = Path("data/raw")
+    OUTPUT_DIR = Path("data/processed/extracted")
+    
+    processor = BatchProcessor(
+        input_dir=INPUT_DIR,
+        output_dir=OUTPUT_DIR,
+        enable_ocr=not args.no_ocr
+    )
+    
+    processor.process_batch(
+        start_year=args.start_year,
+        end_year=args.end_year,
+        limit=args.limit,
+        resume=not args.no_resume
+    )
+
+
+if __name__ == "__main__":
+    main()

src/extraction/pdf_extractor.py

@@ -0,0 +1,522 @@
+"""
+Production PDF Extractor for Legal Judgments
+Enhanced with robust error handling, quality checks, and paragraph preservation
+"""
+
+import PyPDF2
+import pdfplumber
+from pathlib import Path
+from typing import Dict, Optional, List, Tuple
+import logging
+from dataclasses import dataclass, asdict
+import json
+from datetime import datetime
+import re
+
+# OCR imports
+try:
+    import pytesseract
+    from pdf2image import convert_from_path
+    OCR_AVAILABLE = True
+except ImportError:
+    OCR_AVAILABLE = False
+    logging.warning("OCR libraries not installed. OCR fallback disabled.")
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExtractionMetadata:
+    """Metadata for extracted judgment"""
+    filename: str
+    year: str
+    num_pages: int
+    text_length: int
+    extraction_method: str
+    has_text: bool
+    extraction_timestamp: str
+    file_size_bytes: int
+    ocr_used: bool
+    quality_score: float
+    paragraph_count: int
+    errors: List[str]
+    warnings: List[str]
+
+
+class TextQualityChecker:
+    """Utility class for assessing extracted text quality"""
+    
+    # Legal keywords to preserve even in short lines
+    LEGAL_KEYWORDS = {
+        'held', 'order', 'appeal', 'writ', 'judgment', 'decree',
+        'petition', 'application', 'allowed', 'dismissed', 'granted',
+        'rejected', 'reserved', 'disposed', 'quashed', 'set aside',
+        'affirmed', 'reversed', 'remanded', 'suo moto', 'ex parte',
+        'interim', 'stay', 'injunction', 'bail', 'custody', 'liberty',
+        'notice', 'respondent', 'petitioner', 'appellant', 'accused'
+    }
+    
+    @staticmethod
+    def calculate_quality_score(text: str) -> Tuple[float, List[str]]:
+        """
+        Calculate quality score (0-1) for extracted text
+        
+        Returns:
+            (score, issues_found)
+        """
+        if not text or len(text.strip()) < 100:
+            return 0.0, ["Text too short"]
+        
+        issues = []
+        score = 1.0
+        
+        # Check 1: Alphabetic character ratio
+        alpha_chars = sum(c.isalpha() for c in text)
+        total_chars = len(text.replace('\n', '').replace(' ', ''))
+        
+        if total_chars > 0:
+            alpha_ratio = alpha_chars / total_chars
+            if alpha_ratio < 0.5:
+                score -= 0.3
+                issues.append(f"Low alphabetic ratio: {alpha_ratio:.2f}")
+        
+        # Check 2: Average word length (gibberish detection)
+        words = text.split()
+        if words:
+            avg_word_len = sum(len(w) for w in words) / len(words)
+            if avg_word_len < 2 or avg_word_len > 15:
+                score -= 0.2
+                issues.append(f"Unusual avg word length: {avg_word_len:.1f}")
+        
+        # Check 3: Check for repeated patterns (OCR errors)
+        lines = text.split('\n')
+        if len(lines) > 10:
+            unique_lines = len(set(line.strip() for line in lines if line.strip()))
+            repetition_ratio = unique_lines / len(lines)
+            if repetition_ratio < 0.3:
+                score -= 0.2
+                issues.append(f"High repetition: {repetition_ratio:.2f}")
+        
+        # Check 4: Minimum sentence structure
+        sentence_markers = text.count('.') + text.count('?') + text.count('!')
+        if len(words) > 100 and sentence_markers < len(words) / 50:
+            score -= 0.1
+            issues.append("Lacks sentence structure")
+        
+        return max(0.0, min(1.0, score)), issues
+    
+    @staticmethod
+    def clean_ocr_text(text: str) -> str:
+        """
+        Normalize OCR-extracted text with legal-aware filtering
+        - Remove excessive whitespace
+        - Collapse multiple newlines
+        - Remove repeated headers
+        - Preserve important legal terms
+        """
+        # Collapse multiple spaces
+        text = re.sub(r' +', ' ', text)
+        
+        # Collapse multiple newlines (keep max 2 for paragraph breaks)
+        text = re.sub(r'\n{3,}', '\n\n', text)
+        
+        # Remove common OCR artifacts
+        text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f]', '', text)
+        
+        # Legal-aware line filtering
+        lines = text.split('\n')
+        cleaned_lines = []
+        
+        for line in lines:
+            stripped = line.strip()
+            
+            # Skip empty lines
+            if not stripped:
+                continue
+            
+            # Skip pure numbers (page numbers)
+            if stripped.isdigit():
+                continue
+            
+            # PRESERVE if:
+            # 1. Line is substantial (>10 chars)
+            # 2. Contains legal keyword (even if short like "Held.")
+            # 3. Is alphabetic and reasonable length (>3 chars)
+            if (len(stripped) > 10 or
+                any(keyword in stripped.lower() for keyword in TextQualityChecker.LEGAL_KEYWORDS) or
+                (stripped.replace('.', '').replace(',', '').isalpha() and len(stripped) > 3)):
+                cleaned_lines.append(line)
+        
+        text = '\n'.join(cleaned_lines)
+        
+        # Remove repeated header patterns
+        lines = text.split('\n')
+        result = []
+        prev_line = None
+        repeat_count = 0
+        
+        for line in lines:
+            if line.strip() == prev_line and prev_line:
+                repeat_count += 1
+                if repeat_count < 2:  # Allow max 2 repetitions
+                    result.append(line)
+            else:
+                repeat_count = 0
+                result.append(line)
+                prev_line = line.strip()
+        
+        return '\n'.join(result).strip()
+
+
+class LegalJudgmentExtractor:
+    """
+    Production-grade extractor with robust error handling and quality assurance
+    """
+    
+    def __init__(self, output_dir: Path, enable_ocr: bool = True, ocr_max_pages: int = 50):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.enable_ocr = enable_ocr and OCR_AVAILABLE
+        self.ocr_max_pages = ocr_max_pages
+        
+        if enable_ocr and not OCR_AVAILABLE:
+            logger.warning("OCR requested but libraries not installed.")
+        
+        # Create subdirectories
+        self.text_dir = self.output_dir / "texts"
+        self.metadata_dir = self.output_dir / "metadata"
+        self.failed_dir = self.output_dir / "failed"
+        self.ocr_log_file = self.output_dir / "ocr_cases.jsonl"
+        
+        for dir_path in [self.text_dir, self.metadata_dir, self.failed_dir]:
+            dir_path.mkdir(parents=True, exist_ok=True)
+    
+    def extract_year_from_path(self, pdf_path: Path) -> Tuple[str, List[str]]:
+        """
+        Safely extract year from path with validation
+        
+        Returns:
+            (year, warnings)
+        """
+        warnings = []
+        year = pdf_path.parent.name
+        
+        # Validate year
+        if not year.isdigit():
+            warnings.append(f"Invalid year from directory: {year}")
+            
+            # Try to extract from filename
+            filename = pdf_path.stem
+            year_match = re.search(r'(19|20)\d{2}', filename)
+            if year_match:
+                year = year_match.group(0)
+                warnings.append(f"Year extracted from filename: {year}")
+            else:
+                year = "unknown"
+                warnings.append("Could not determine year")
+        else:
+            # Validate year range
+            year_int = int(year)
+            if year_int < 1950 or year_int > 2025:
+                warnings.append(f"Year {year} outside expected range (1950-2025)")
+        
+        return year, warnings
+    
+    def count_paragraphs(self, text: str) -> int:
+        """Count paragraph-like structures in text"""
+        # Split by double newlines
+        paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
+        # Filter out very short "paragraphs" (likely headers)
+        substantial_paragraphs = [p for p in paragraphs if len(p) > 50]
+        return len(substantial_paragraphs)
+    
+    def extract_with_pypdf2(self, pdf_path: Path) -> Optional[str]:
+        """Primary extraction - preserves paragraph structure"""
+        try:
+            with open(pdf_path, 'rb') as file:
+                reader = PyPDF2.PdfReader(file)
+                text_parts = []
+                
+                for page in reader.pages:
+                    text = page.extract_text()
+                    if text:
+                        text_parts.append(text.strip())
+                
+                # Join with double newline to preserve page breaks
+                full_text = "\n\n".join(text_parts)
+                
+                # Quality check
+                score, _ = TextQualityChecker.calculate_quality_score(full_text)
+                return full_text if score > 0.3 else None
+                
+        except Exception as e:
+            logger.debug(f"PyPDF2 failed for {pdf_path.name}: {e}")
+            return None
+    
+    def extract_with_pdfplumber(self, pdf_path: Path) -> Optional[str]:
+        """Fallback extraction - better for complex layouts"""
+        try:
+            with pdfplumber.open(pdf_path) as pdf:
+                text_parts = []
+                
+                for page in pdf.pages:
+                    text = page.extract_text()
+                    if text:
+                        text_parts.append(text.strip())
+                
+                full_text = "\n\n".join(text_parts)
+                
+                score, _ = TextQualityChecker.calculate_quality_score(full_text)
+                return full_text if score > 0.3 else None
+                
+        except Exception as e:
+            logger.debug(f"pdfplumber failed for {pdf_path.name}: {e}")
+            return None
+    
+    def extract_with_ocr(self, pdf_path: Path, num_pages: int) -> Optional[str]:
+        """
+        OCR extraction with proper page limiting and text normalization
+        
+        Args:
+            pdf_path: Path to PDF
+            num_pages: Total pages in PDF (for proper limiting)
+        """
+        if not self.enable_ocr:
+            return None
+        
+        try:
+            logger.info(f"OCR extraction: {pdf_path.name}")
+            
+            # Proper page limiting
+            last_page = min(self.ocr_max_pages, num_pages)
+            
+            if num_pages > self.ocr_max_pages:
+                logger.warning(f"PDF has {num_pages} pages, OCR limited to first {self.ocr_max_pages}")
+            
+            # Convert to images
+            images = convert_from_path(
+                pdf_path,
+                dpi=300,
+                first_page=1,
+                last_page=last_page
+            )
+            
+            text_parts = []
+            for i, image in enumerate(images, 1):
+                logger.debug(f"OCR page {i}/{len(images)}")
+                
+                text = pytesseract.image_to_string(image, lang='eng')
+                if text.strip():
+                    text_parts.append(text)
+            
+            full_text = "\n\n".join(text_parts)
+            
+            # Normalize OCR text
+            full_text = TextQualityChecker.clean_ocr_text(full_text)
+            
+            # Check quality
+            score, issues = TextQualityChecker.calculate_quality_score(full_text)
+            
+            if score > 0.3:
+                # Log successful OCR to JSONL
+                self._log_ocr_case(pdf_path, num_pages, last_page, score)
+                logger.info(f"✓ OCR successful (quality: {score:.2f})")
+                return full_text
+            else:
+                logger.warning(f"OCR quality too low ({score:.2f}): {issues}")
+                return None
+                
+        except Exception as e:
+            logger.warning(f"OCR failed for {pdf_path.name}: {e}")
+            return None
+    
+    def _log_ocr_case(self, pdf_path: Path, total_pages: int, pages_processed: int, quality: float):
+        """Log OCR usage to JSONL file"""
+        log_entry = {
+            'timestamp': datetime.now().isoformat(),
+            'filename': pdf_path.name,
+            'year': pdf_path.parent.name,
+            'total_pages': total_pages,
+            'pages_processed': pages_processed,
+            'quality_score': quality
+        }
+        
+        with open(self.ocr_log_file, 'a', encoding='utf-8') as f:
+            f.write(json.dumps(log_entry) + '\n')
+    
+    def extract_pdf(self, pdf_path: Path) -> Dict:
+        """
+        Main extraction with fallback chain and quality assurance
+        """
+        errors = []
+        warnings = []
+        text = None
+        method = None
+        ocr_used = False
+        quality_score = 0.0
+        
+        # Get metadata
+        file_size = pdf_path.stat().st_size
+        
+        # Robust year extraction
+        year, year_warnings = self.extract_year_from_path(pdf_path)
+        warnings.extend(year_warnings)
+        
+        # Count pages first (needed for OCR)
+        try:
+            with open(pdf_path, 'rb') as f:
+                reader = PyPDF2.PdfReader(f)
+                num_pages = len(reader.pages)
+        except Exception as e:
+            num_pages = 0
+            errors.append(f"Could not count pages: {e}")
+        
+        # Extraction chain: PyPDF2 → pdfplumber → OCR
+        text = self.extract_with_pypdf2(pdf_path)
+        if text:
+            method = "pypdf2"
+        else:
+            errors.append("PyPDF2 insufficient")
+            
+            text = self.extract_with_pdfplumber(pdf_path)
+            if text:
+                method = "pdfplumber"
+            else:
+                errors.append("pdfplumber failed")
+                
+                if self.enable_ocr and num_pages > 0:
+                    text = self.extract_with_ocr(pdf_path, num_pages)
+                    if text:
+                        method = "ocr"
+                        ocr_used = True
+                        warnings.append("OCR used - verify quality")
+                    else:
+                        errors.append("OCR failed")
+        
+        # Calculate quality
+        paragraph_count = 0
+        if text:
+            quality_score, quality_issues = TextQualityChecker.calculate_quality_score(text)
+            paragraph_count = self.count_paragraphs(text)
+            
+            if quality_score < 0.7:
+                warnings.extend(quality_issues)
+        
+        # Create metadata
+        metadata = ExtractionMetadata(
+            filename=pdf_path.name,
+            year=year,
+            num_pages=num_pages,
+            text_length=len(text) if text else 0,
+            extraction_method=method if method else "failed",
+            has_text=text is not None,
+            extraction_timestamp=datetime.now().isoformat(),
+            file_size_bytes=file_size,
+            ocr_used=ocr_used,
+            quality_score=quality_score,
+            paragraph_count=paragraph_count,
+            errors=errors,
+            warnings=warnings
+        )
+        
+        return {
+            'text': text,
+            'metadata': metadata
+        }
+    
+    def save_extraction(self, pdf_path: Path, extraction_result: Dict) -> bool:
+        """Save with quality indicators"""
+        
+        metadata = extraction_result['metadata']
+        text = extraction_result['text']
+        
+        base_name = pdf_path.stem
+        year = metadata.year
+        
+        # Save text
+        if text:
+            text_file = self.text_dir / f"{year}_{base_name}.txt"
+            try:
+                with open(text_file, 'w', encoding='utf-8') as f:
+                    # Add quality header
+                    f.write(f"{'='*70}\n")
+                    f.write(f"File: {metadata.filename}\n")
+                    f.write(f"Extraction: {metadata.extraction_method}\n")
+                    f.write(f"Quality: {metadata.quality_score:.2f}\n")
+                    f.write(f"Paragraphs: {metadata.paragraph_count}\n")
+                    
+                    if metadata.ocr_used:
+                        f.write("⚠️ OCR USED - Verify important details\n")
+                    
+                    if metadata.warnings:
+                        f.write(f"Warnings: {', '.join(metadata.warnings[:3])}\n")
+                    
+                    f.write(f"{'='*70}\n\n")
+                    f.write(text)
+                    
+            except Exception as e:
+                logger.error(f"Failed to save text: {e}")
+                return False
+        
+        # Save metadata
+        metadata_file = self.metadata_dir / f"{year}_{base_name}.json"
+        try:
+            with open(metadata_file, 'w', encoding='utf-8') as f:
+                json.dump(asdict(metadata), f, indent=2)
+        except Exception as e:
+            logger.error(f"Failed to save metadata: {e}")
+            return False
+        
+        # Log failures
+        if not text:
+            failed_log = self.failed_dir / "failed_extractions.jsonl"
+            with open(failed_log, 'a', encoding='utf-8') as f:
+                log_entry = {
+                    'timestamp': datetime.now().isoformat(),
+                    'file': str(pdf_path),
+                    'errors': metadata.errors
+                }
+                f.write(json.dumps(log_entry) + '\n')
+        
+        return True
+    
+    def process_pdf(self, pdf_path: Path) -> bool:
+        """Process single PDF"""
+        try:
+            result = self.extract_pdf(pdf_path)
+            return self.save_extraction(pdf_path, result)
+        except Exception as e:
+            logger.error(f"Unexpected error: {pdf_path.name}: {e}")
+            return False
+
+
+if __name__ == "__main__":
+    # Test
+    print("="*70)
+    print("Testing Enhanced PDF Extractor")
+    print("="*70)
+    
+    extractor = LegalJudgmentExtractor(
+        output_dir=Path("data/processed/extracted"),
+        enable_ocr=False
+    )
+    
+    test_pdf = Path("data/raw/2025/A_John_Kennedy_vs_The_State_Of_Tamil_Nadu_on_24_March_2025_1.PDF")
+    
+    if test_pdf.exists():
+        print(f"\nTesting: {test_pdf.name}")
+        success = extractor.process_pdf(test_pdf)
+        print(f"\n{'✓' if success else '✗'} Extraction {'successful' if success else 'failed'}")
+        
+        # Show metadata
+        metadata_file = Path("data/processed/extracted/metadata") / f"2025_{test_pdf.stem}.json"
+        if metadata_file.exists():
+            with open(metadata_file, 'r') as f:
+                metadata = json.load(f)
+                print(f"\nMethod: {metadata['extraction_method']}")
+                print(f"Quality: {metadata['quality_score']:.2f}")
+                print(f"Paragraphs: {metadata['paragraph_count']}")
+                print(f"Text length: {metadata['text_length']:,} chars")
+    else:
+        print("Test PDF not found")

src/indexing/build_faiss_index.py

@@ -0,0 +1,85 @@
+"""
+NyayLens – FAISS Index Builder (Merged Embeddings)
+"""
+
+import faiss
+import numpy as np
+import json
+from pathlib import Path
+
+
+def build_faiss_index():
+    print("=" * 70)
+    print("NyayLens – Building FAISS Index")
+    print("=" * 70)
+
+    embeddings_dir = Path("data/processed/embeddings")
+    output_dir = Path("data/processed/faiss")
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    embeddings_file = embeddings_dir / "paragraph_embeddings.npy"
+    ids_file = embeddings_dir / "paragraph_ids.json"
+    meta_file = embeddings_dir / "embedding_metadata.json"
+
+    # --- Load data ---
+    print("\nLoading embeddings...")
+    embeddings = np.load(embeddings_file)
+
+    print("Loading paragraph IDs...")
+    with open(ids_file, "r") as f:
+        paragraph_ids = json.load(f)
+
+    # --- Safety checks ---
+    assert embeddings.shape[0] == len(paragraph_ids), (
+        f"Mismatch: {embeddings.shape[0]} embeddings vs "
+        f"{len(paragraph_ids)} paragraph IDs"
+    )
+
+    print(f"✓ Loaded {embeddings.shape[0]:,} vectors")
+    print(f"✓ Embedding dimension: {embeddings.shape[1]}")
+
+    # --- Normalize (cosine similarity) ---
+    print("\nNormalizing embeddings...")
+    faiss.normalize_L2(embeddings)
+
+    # --- Build FAISS index ---
+    dim = embeddings.shape[1]
+    
+    # Switch to HNSW for rapid Approximate Nearest Neighbor search
+    # M=32 is number of connections per layer, typical for dense models
+    index = faiss.IndexHNSWFlat(dim, 32, faiss.METRIC_INNER_PRODUCT)
+    index.hnsw.efConstruction = 40  # Depth of search during index build
+    index.hnsw.efSearch = 64        # Depth of search during inference
+
+    print(f"Adding {embeddings.shape[0]:,} vectors to HNSW index...")
+    index.add(embeddings)
+
+    print(f"✓ FAISS index contains {index.ntotal:,} vectors")
+
+    # --- Save index ---
+    index_file = output_dir / "faiss_index.bin"
+    faiss.write_index(index, str(index_file))
+
+    # --- Save FAISS metadata ---
+    with open(meta_file, "r") as f:
+        emb_meta = json.load(f)
+
+    faiss_meta = {
+        "index_type": "IndexHNSWFlat",
+        "metric": "cosine_similarity",
+        "dimension": dim,
+        "total_vectors": int(index.ntotal),
+        "embedding_model": emb_meta.get("model_name", "unknown"),
+    }
+
+    with open(output_dir / "faiss_metadata.json", "w") as f:
+        json.dump(faiss_meta, f, indent=2)
+
+    print(f"\n✓ FAISS index saved to: {index_file}")
+    print("=" * 70)
+    print("✓ FAISS Index Build Complete")
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    build_faiss_index()

src/indexing/create_embeddings.py

@@ -0,0 +1,232 @@
+"""
+Generate embeddings in CHUNKS - GPU-safe with resume capability
+Processes 20K paragraphs at a time with cooling breaks
+"""
+
+from sentence_transformers import SentenceTransformer
+import json
+import numpy as np
+from pathlib import Path
+from tqdm import tqdm
+import torch
+import time
+
+class ChunkedEmbeddingGenerator:
+    """Generate embeddings in safe chunks with resume capability"""
+    
+    def __init__(self, model_name: str = "BAAI/bge-base-en-v1.5", chunk_size: int = 20000):
+        """
+        Initialize with sentence transformer
+        chunk_size: Process this many paragraphs at a time
+        """
+        print(f"Loading model: {model_name}")
+        
+        # Check CUDA availability
+        if torch.cuda.is_available():
+            print(f"✓ CUDA available! GPU: {torch.cuda.get_device_name(0)}")
+            print(f"  GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
+            self.device = 'cuda'
+        else:
+            print("⚠️  CUDA not available, using CPU")
+            self.device = 'cpu'
+        
+        self.model = SentenceTransformer(model_name)
+        self.model.to(self.device)
+        self.chunk_size = chunk_size
+        print(f"✓ Model loaded on: {self.device}")
+        print(f"✓ Chunk size: {chunk_size:,} paragraphs")
+    
+    def load_paragraphs(self, index_file: Path):
+        """Load paragraphs from JSONL"""
+        print("\nLoading paragraphs from index...")
+        paragraphs = []
+        
+        with open(index_file, 'r', encoding='utf-8') as f:
+            total_lines = sum(1 for _ in f)
+        
+        with open(index_file, 'r', encoding='utf-8') as f:
+            for line in tqdm(f, total=total_lines, desc="Loading index"):
+                para = json.loads(line)
+                paragraphs.append(para)
+        
+        print(f"✓ Loaded {len(paragraphs):,} paragraphs")
+        return paragraphs
+    
+    def process_chunk(self, texts, batch_size=64):
+        """Process one chunk of texts"""
+        embeddings_list = []
+        
+        with tqdm(total=len(texts), desc="Processing chunk", unit="para") as pbar:
+            for i in range(0, len(texts), batch_size):
+                batch = texts[i:i+batch_size]
+                
+                # Generate embeddings
+                batch_embeddings = self.model.encode(
+                    batch,
+                    convert_to_numpy=True,
+                    normalize_embeddings=True,
+                    show_progress_bar=False,
+                    device=self.device
+                )
+                
+                embeddings_list.append(batch_embeddings)
+                pbar.update(len(batch))
+                
+                # Clear GPU cache every 10 batches
+                if self.device == 'cuda' and i % (batch_size * 10) == 0:
+                    torch.cuda.empty_cache()
+        
+        return np.vstack(embeddings_list)
+    
+    def generate_embeddings_chunked(self, index_file: Path, output_dir: Path, batch_size: int = 64):
+        """Generate embeddings in chunks with resume capability"""
+        
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        chunks_dir = output_dir / "chunks"
+        chunks_dir.mkdir(exist_ok=True)
+        
+        # Load paragraphs
+        paragraphs = self.load_paragraphs(index_file)
+        total_paras = len(paragraphs)
+        
+        # Calculate chunks
+        num_chunks = (total_paras + self.chunk_size - 1) // self.chunk_size
+        print(f"\n✓ Will process in {num_chunks} chunks of {self.chunk_size:,} paragraphs each")
+        
+        # Check for existing chunks
+        existing_chunks = list(chunks_dir.glob("chunk_*.npy"))
+        completed_chunks = len(existing_chunks)
+        
+        if completed_chunks > 0:
+            print(f"✓ Found {completed_chunks} existing chunks")
+            response = input(f"Resume from chunk {completed_chunks + 1}? (yes/no): ")
+            if response.lower() != 'yes':
+                print("Starting fresh...")
+                for f in existing_chunks:
+                    f.unlink()
+                completed_chunks = 0
+        
+        # Process chunks
+        start_time = time.time()
+        
+        for chunk_idx in range(completed_chunks, num_chunks):
+            print(f"\n{'='*70}")
+            print(f"CHUNK {chunk_idx + 1}/{num_chunks}")
+            print(f"{'='*70}")
+            
+            # Get chunk data
+            start_idx = chunk_idx * self.chunk_size
+            end_idx = min(start_idx + self.chunk_size, total_paras)
+            chunk_paras = paragraphs[start_idx:end_idx]
+            
+            print(f"Processing paragraphs {start_idx:,} to {end_idx:,}")
+            
+            # Extract texts
+            texts = [p['text'] for p in chunk_paras]
+            ids = [p['id'] for p in chunk_paras]
+            
+            # Process chunk
+            chunk_start = time.time()
+            embeddings = self.process_chunk(texts, batch_size)
+            chunk_time = time.time() - chunk_start
+            
+            print(f"✓ Chunk completed in {chunk_time/60:.1f} minutes")
+            print(f"  Speed: {len(texts)/chunk_time:.1f} para/s")
+            
+            # Save chunk
+            chunk_file = chunks_dir / f"chunk_{chunk_idx:03d}.npy"
+            ids_file = chunks_dir / f"chunk_{chunk_idx:03d}_ids.json"
+            
+            np.save(chunk_file, embeddings)
+            with open(ids_file, 'w') as f:
+                json.dump(ids, f)
+            
+            print(f"✓ Saved to: {chunk_file.name}")
+            
+            # Clear GPU and add cooling break
+            if self.device == 'cuda':
+                torch.cuda.empty_cache()
+                if chunk_idx < num_chunks - 1:  # Not last chunk
+                    print("\n⏸️  Cooling break: 10 seconds...")
+                    time.sleep(10)
+        
+        # Combine all chunks
+        print(f"\n{'='*70}")
+        print("COMBINING CHUNKS...")
+        print(f"{'='*70}")
+        
+        all_embeddings = []
+        all_ids = []
+        
+        for chunk_idx in tqdm(range(num_chunks), desc="Loading chunks"):
+            chunk_file = chunks_dir / f"chunk_{chunk_idx:03d}.npy"
+            ids_file = chunks_dir / f"chunk_{chunk_idx:03d}_ids.json"
+            
+            embeddings = np.load(chunk_file)
+            with open(ids_file, 'r') as f:
+                ids = json.load(f)
+            
+            all_embeddings.append(embeddings)
+            all_ids.extend(ids)
+        
+        final_embeddings = np.vstack(all_embeddings)
+        
+        print(f"✓ Combined shape: {final_embeddings.shape}")
+        
+        # Save final files
+        print("\nSaving final files...")
+        
+        embeddings_file = output_dir / "paragraph_embeddings.npy"
+        np.save(embeddings_file, final_embeddings)
+        print(f"✓ Saved: {embeddings_file}")
+        
+        ids_file = output_dir / "paragraph_ids.json"
+        with open(ids_file, 'w') as f:
+            json.dump(all_ids, f)
+        print(f"✓ Saved: {ids_file}")
+        
+        # Save metadata
+        total_time = time.time() - start_time
+        metadata = {
+            'model_name': "BAAI/bge-base-en-v1.5",
+            'embedding_dim': int(final_embeddings.shape[1]),
+            'total_paragraphs': len(all_ids),
+            'device_used': self.device,
+            'batch_size': batch_size,
+            'chunk_size': self.chunk_size,
+            'num_chunks': num_chunks,
+            'total_time_minutes': total_time / 60
+        }
+        
+        with open(output_dir / "embedding_metadata.json", 'w') as f:
+            json.dump(metadata, f, indent=2)
+        
+        print(f"\n✓ Total time: {total_time/60:.1f} minutes")
+        print(f"  Average speed: {len(all_ids)/total_time:.1f} para/s")
+        
+        return final_embeddings, all_ids
+
+
+if __name__ == "__main__":
+    print("="*70)
+    print("NyayLens - Chunked Embedding Generation (GPU-Safe)")
+    print("="*70)
+    print()
+    
+    generator = ChunkedEmbeddingGenerator(
+        model_name="BAAI/bge-base-en-v1.5",
+        chunk_size=20000  # Process 20K at a time
+    )
+    
+    embeddings, ids = generator.generate_embeddings_chunked(
+        index_file=Path("data/processed/indexed/paragraph_index.jsonl"),
+        output_dir=Path("data/processed/embeddings"),
+        batch_size=64  # Reduced for safety
+    )
+    
+    print()
+    print("="*70)
+    print(f"✓ COMPLETE! Generated {len(embeddings):,} embeddings")
+    print("="*70)

src/indexing/create_sqlite_index.py

@@ -0,0 +1,196 @@
+# """Create SQLite index for fast paragraph lookup"""
+
+# import sqlite3
+# import json
+# from pathlib import Path
+# from tqdm import tqdm
+
+# def create_sqlite_index():
+#     print("Creating SQLite index...")
+    
+#     db_path = Path("data/processed/indexed/paragraphs.db")
+#     db_path.parent.mkdir(parents=True, exist_ok=True)
+    
+#     # Create database
+#     conn = sqlite3.connect(db_path)
+#     cursor = conn.cursor()
+    
+#     # Create table
+#     cursor.execute("""
+#         CREATE TABLE IF NOT EXISTS paragraphs (
+#             id TEXT PRIMARY KEY,
+#             judgment_id TEXT,
+#             page_no INTEGER,
+#             text TEXT,
+#             char_count INTEGER,
+#             word_count INTEGER
+#         )
+#     """)
+    
+#     cursor.execute("CREATE INDEX IF NOT EXISTS idx_judgment ON paragraphs(judgment_id)")
+    
+#     # Load data
+#     index_file = Path("data/processed/indexed/paragraph_index.jsonl")
+    
+#     with open(index_file, 'r', encoding='utf-8') as f:
+#         total = sum(1 for _ in f)
+    
+#     with open(index_file, 'r', encoding='utf-8') as f:
+#         batch = []
+#         for line in tqdm(f, total=total, desc="Inserting"):
+#             p = json.loads(line)
+#             batch.append((
+#                 p['id'], p['judgment_id'], p['page_no'],
+#                 p['text'], p['char_count'], p['word_count']
+#             ))
+            
+#             if len(batch) >= 1000:
+#                 cursor.executemany(
+#                     "INSERT OR REPLACE INTO paragraphs VALUES (?,?,?,?,?,?)",
+#                     batch
+#                 )
+#                 batch = []
+        
+#         if batch:
+#             cursor.executemany(
+#                 "INSERT OR REPLACE INTO paragraphs VALUES (?,?,?,?,?,?)",
+#                 batch
+#             )
+    
+#     conn.commit()
+#     conn.close()
+    
+#     print(f"✓ SQLite index created: {db_path}")
+
+# if __name__ == "__main__":
+#     create_sqlite_index()
+"""
+Create SQLite index with section annotations
+Source: paragraph_index_with_sections.jsonl
+"""
+
+import sqlite3
+import json
+from pathlib import Path
+from tqdm import tqdm
+
+
+INPUT_INDEX = Path("data/processed/indexed/paragraph_index_with_sections.jsonl")
+DB_PATH = Path("data/processed/indexed/paragraphs.db")
+
+
+def create_sqlite_index():
+    print("=" * 70)
+    print("NyayLens – Creating SQLite Index (with Sections)")
+    print("=" * 70)
+
+    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+
+    # Connect to SQLite
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+
+    # Drop existing table (derived data → safe to rebuild)
+    cursor.execute("DROP TABLE IF EXISTS paragraphs")
+
+    # Create table
+    cursor.execute("""
+        CREATE TABLE paragraphs (
+            id TEXT PRIMARY KEY,
+            judgment_id TEXT,
+            page_no INTEGER,
+            text TEXT,
+            char_count INTEGER,
+            word_count INTEGER,
+            section TEXT,
+            section_conf REAL
+        )
+    """)
+
+    # Create FTS5 virtual table for fast full-text search (BM25)
+    cursor.execute("DROP TABLE IF EXISTS paragraphs_fts")
+    cursor.execute("""
+        CREATE VIRTUAL TABLE paragraphs_fts USING fts5(
+            id UNINDEXED,
+            text,
+            tokenize='porter unicode61'
+        )
+    """)
+
+    # Indexes for fast lookup
+    cursor.execute("CREATE INDEX idx_judgment_id ON paragraphs(judgment_id)")
+    cursor.execute("CREATE INDEX idx_section ON paragraphs(section)")
+    cursor.execute("CREATE INDEX idx_judgment_section ON paragraphs(judgment_id, section)")
+
+    conn.commit()
+
+    # Count total records
+    with open(INPUT_INDEX, "r", encoding="utf-8") as f:
+        total = sum(1 for _ in f)
+
+    print(f"✓ Inserting {total:,} paragraphs")
+
+    # Insert data in batches
+    batch = []
+    BATCH_SIZE = 1000
+
+    with open(INPUT_INDEX, "r", encoding="utf-8") as f:
+        for line in tqdm(f, total=total, desc="Inserting"):
+            p = json.loads(line)
+
+            batch.append((
+                p["id"],
+                p["judgment_id"],
+                p.get("page_no", -1),
+                p["text"],
+                p.get("char_count", len(p["text"])),
+                p.get("word_count", len(p["text"].split())),
+                p.get("section", "unknown"),
+                p.get("section_conf", 0.0),
+            ))
+
+            if len(batch) >= BATCH_SIZE:
+                cursor.executemany(
+                    """
+                    INSERT INTO paragraphs
+                    (id, judgment_id, page_no, text, char_count, word_count, section, section_conf)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                    batch
+                )
+                
+                # Insert into FTS5 table
+                fts_batch = [(b[0], b[3]) for b in batch]
+                cursor.executemany(
+                    "INSERT INTO paragraphs_fts (id, text) VALUES (?, ?)",
+                    fts_batch
+                )
+                
+                batch.clear()
+
+        if batch:
+            cursor.executemany(
+                """
+                INSERT INTO paragraphs
+                (id, judgment_id, page_no, text, char_count, word_count, section, section_conf)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                batch
+            )
+            
+            fts_batch = [(b[0], b[3]) for b in batch]
+            cursor.executemany(
+                "INSERT INTO paragraphs_fts (id, text) VALUES (?, ?)",
+                fts_batch
+            )
+
+    conn.commit()
+    conn.close()
+
+    print("\n✓ SQLite index created successfully")
+    print(f"✓ Database path: {DB_PATH}")
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    create_sqlite_index()

src/indexing/paragraph_indexer.py

@@ -0,0 +1,167 @@
+"""
+Paragraph-level indexer for NyayLens RAG
+- Page-aware
+- Content-stable paragraph IDs
+- Legal-aware filtering
+- Streaming JSONL output (memory safe)
+"""
+
+import json
+import hashlib
+import re
+from pathlib import Path
+from typing import Dict, Iterable
+from tqdm import tqdm
+
+
+class ParagraphIndexer:
+    """Index legal judgments at paragraph level with stable IDs and metadata"""
+
+    # Legal keywords worth preserving even in short paragraphs
+    LEGAL_KEYWORDS = {
+        'held', 'order', 'appeal', 'writ', 'judgment', 'decree',
+        'petition', 'application', 'allowed', 'dismissed', 'granted',
+        'rejected', 'disposed', 'quashed', 'set aside',
+        'affirmed', 'reversed', 'remanded', 'bail', 'custody',
+        'interim', 'stay', 'injunction', 'no costs'
+    }
+
+    PAGE_MARKER_PATTERN = re.compile(r'<<<PAGE:(\d+)>>>')
+
+    def __init__(self, texts_dir: Path, output_dir: Path):
+        self.texts_dir = Path(texts_dir)
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        self.index_file = self.output_dir / "paragraph_index.jsonl"
+        self.stats_file = self.output_dir / "index_stats.json"
+
+    @staticmethod
+    def _contains_legal_keyword(text: str) -> bool:
+        text_l = text.lower()
+        return any(
+            re.search(rf"\b{re.escape(kw)}\b", text_l)
+            for kw in ParagraphIndexer.LEGAL_KEYWORDS
+        )
+
+    @staticmethod
+    def _stable_paragraph_id(judgment_id: str, page_no: int, text: str) -> str:
+        """
+        Content-stable ID:
+        - same paragraph text => same ID
+        - survives re-indexing
+        """
+        h = hashlib.sha1(text.encode("utf-8")).hexdigest()[:16]
+        page_str = page_no if page_no is not None else "unk"
+        return f"{judgment_id}_p{page_str}_{h}"
+
+    def _strip_header(self, content: str) -> str:
+        """
+        Remove extractor quality header safely
+        """
+        sep = "=" * 70
+        if sep in content:
+            parts = content.split(sep, 2)
+            if len(parts) == 3:
+                return parts[2].strip()
+        return content.strip()
+
+    def _iter_paragraphs(self, content: str) -> Iterable[tuple]:
+        """
+        Yield paragraph records with page numbers
+        """
+        current_page = None
+        buffer = []
+
+        for line in content.splitlines():
+            page_match = self.PAGE_MARKER_PATTERN.match(line.strip())
+            if page_match:
+                # Flush buffer before page change
+                if buffer:
+                    yield current_page, "\n".join(buffer).strip()
+                    buffer = []
+                current_page = int(page_match.group(1))
+                continue
+
+            if not line.strip():
+                if buffer:
+                    yield current_page, "\n".join(buffer).strip()
+                    buffer = []
+                continue
+
+            buffer.append(line)
+
+        if buffer:
+            yield current_page, "\n".join(buffer).strip()
+
+    def index_judgment(self, text_file: Path, writer) -> int:
+        """
+        Index a single judgment file.
+        Returns number of paragraphs indexed.
+        """
+        with open(text_file, "r", encoding="utf-8") as f:
+            content = self._strip_header(f.read())
+
+        judgment_id = text_file.stem
+        para_count = 0
+
+        for page_no, para in self._iter_paragraphs(content):
+            if not para:
+                continue
+
+            # Keep substantial OR legally important short paragraphs
+            if len(para) < 50 and not self._contains_legal_keyword(para):
+                continue
+
+            record = {
+                "id": self._stable_paragraph_id(judgment_id, page_no if page_no is not None else -1, para),
+                "judgment_id": judgment_id,
+                "page_no": page_no if page_no is not None else -1,
+                "text": para,
+                "char_count": len(para),
+                "word_count": len(para.split())
+            }
+
+            writer.write(json.dumps(record, ensure_ascii=False) + "\n")
+            para_count += 1
+
+        return para_count
+
+    def build_full_index(self):
+        text_files = sorted(self.texts_dir.glob("*.txt"))
+        print(f"Indexing {len(text_files):,} judgments...")
+
+        total_paragraphs = 0
+
+        with open(self.index_file, "w", encoding="utf-8") as writer:
+            for text_file in tqdm(text_files, desc="Indexing"):
+                try:
+                    total_paragraphs += self.index_judgment(text_file, writer)
+                except Exception as e:
+                    print(f"❌ Failed indexing {text_file.name}: {e}")
+
+        stats = {
+            "total_judgments": len(text_files),
+            "total_paragraphs": total_paragraphs,
+            "avg_paragraphs_per_judgment":
+                total_paragraphs / len(text_files) if text_files else 0
+        }
+
+        with open(self.stats_file, "w", encoding="utf-8") as f:
+            json.dump(stats, f, indent=2)
+
+        print("\n✓ Paragraph indexing complete")
+        print(f"  Total paragraphs: {total_paragraphs:,}")
+        print(f"  Output: {self.index_file}")
+
+        return stats
+
+
+if __name__ == "__main__":
+    indexer = ParagraphIndexer(
+        texts_dir=Path("data/processed/extracted/texts"),
+        output_dir=Path("data/processed/indexed")
+    )
+
+    stats = indexer.build_full_index()
+    print(f"\nAverage paragraphs per judgment: {stats['avg_paragraphs_per_judgment']:.1f}")

src/pipeline.py

@@ -0,0 +1,84 @@
+import argparse
+from pathlib import Path
+import logging
+import time
+
+from extraction.pdf_extractor import LegalJudgmentExtractor
+from segmentation.judgement_segmenter import JudgmentSegmenter
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class NyayLensPipeline:
+    """Unified ingestion pipeline for NyayLens"""
+    
+    def __init__(self, raw_dir="data/raw", processed_dir="data/processed"):
+        self.raw_dir = Path(raw_dir)
+        self.processed_dir = Path(processed_dir)
+        
+        logger.info("Initializing NyayLens Pipeline components...")
+        self.extractor = LegalJudgmentExtractor(output_dir=self.processed_dir / "extracted")
+        self.segmenter = JudgmentSegmenter()
+        
+    def ingest_pdf(self, pdf_path: Path):
+        """Process a single PDF end-to-end"""
+        logger.info(f"--- Starting Ingestion: {pdf_path.name} ---")
+        start_time = time.time()
+        
+        # 1. Extraction
+        logger.info("Step 1: Extracting text...")
+        extraction_result = self.extractor.extract_pdf(pdf_path)
+        if not extraction_result['text']:
+            logger.error(f"Extraction failed for {pdf_path.name}")
+            return False
+            
+        # Optional: Save extraction
+        self.extractor.save_extraction(pdf_path, extraction_result)
+        
+        # 2. Segmentation
+        logger.info("Step 2: Segmenting judgment...")
+        # Simple paragraph split for segmentation
+        paragraphs = extraction_result['text'].split('\n\n')
+        sections = self.segmenter.segment(paragraphs)
+        
+        logger.info(f"Found {len(sections)} distinct sections.")
+        
+        # 3. Next steps would hook into `create_embeddings.py` and `create_sqlite_index.py`
+        logger.info("Step 3: Ready for Embeddings & Indexing (Batch process recommended)")
+        
+        elapsed = time.time() - start_time
+        logger.info(f"--- Successfully processed {pdf_path.name} in {elapsed:.2f}s ---")
+        return True
+
+    def process_directory(self, limit: int = None):
+        """Process all PDFs in the raw directory"""
+        pdfs = list(self.raw_dir.glob("**/*.pdf")) + list(self.raw_dir.glob("**/*.PDF"))
+        
+        if limit:
+            pdfs = pdfs[:limit]
+            
+        logger.info(f"Found {len(pdfs)} PDFs to process.")
+        
+        success = 0
+        for pdf in pdfs:
+            if self.ingest_pdf(pdf):
+                success += 1
+                
+        logger.info(f"Batch completed. {success}/{len(pdfs)} successful.")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="NyayLens Unified Ingestion Pipeline")
+    parser.add_argument("--pdf", type=str, help="Path to a single PDF to ingest")
+    parser.add_argument("--batch", action="store_true", help="Process all PDFs in data/raw")
+    parser.add_argument("--limit", type=int, default=None, help="Limit number of PDFs in batch")
+    
+    args = parser.parse_args()
+    
+    pipeline = NyayLensPipeline()
+    
+    if args.pdf:
+        pipeline.ingest_pdf(Path(args.pdf))
+    elif args.batch:
+        pipeline.process_directory(limit=args.limit)
+    else:
+        logger.warning("Please specify --pdf <path> or --batch. Use --help for options.")

src/qa/dataset.py

@@ -0,0 +1,69 @@
+from datasets import load_dataset
+
+MAX_LENGTH = 384
+DOC_STRIDE = 128
+
+def load_and_prepare_dataset(tokenizer):
+    dataset = load_dataset("squad")  # auto-download
+
+    def preprocess(examples):
+        questions = [q.strip() for q in examples["question"]]
+        contexts = examples["context"]
+
+        tokenized = tokenizer(
+            questions,
+            contexts,
+            truncation="only_second",
+            max_length=MAX_LENGTH,
+            stride=DOC_STRIDE,
+            return_overflowing_tokens=True,
+            return_offsets_mapping=True,
+            padding="max_length",
+        )
+
+        sample_mapping = tokenized.pop("overflow_to_sample_mapping")
+        offset_mapping = tokenized.pop("offset_mapping")
+
+        start_positions = []
+        end_positions = []
+
+        for i, offsets in enumerate(offset_mapping):
+            input_ids = tokenized["input_ids"][i]
+            cls_index = input_ids.index(tokenizer.cls_token_id)
+
+            sample_idx = sample_mapping[i]
+            answer = examples["answers"][sample_idx]
+
+            if len(answer["answer_start"]) == 0:
+                start_positions.append(cls_index)
+                end_positions.append(cls_index)
+            else:
+                start_char = answer["answer_start"][0]
+                end_char = start_char + len(answer["text"][0])
+
+                token_start = token_end = None
+                for idx, (start, end) in enumerate(offsets):
+                    if start <= start_char < end:
+                        token_start = idx
+                    if start < end_char <= end:
+                        token_end = idx
+                        break
+
+                if token_start is None or token_end is None:
+                    start_positions.append(cls_index)
+                    end_positions.append(cls_index)
+                else:
+                    start_positions.append(token_start)
+                    end_positions.append(token_end)
+
+        tokenized["start_positions"] = start_positions
+        tokenized["end_positions"] = end_positions
+        return tokenized
+
+    tokenized_dataset = dataset.map(
+        preprocess,
+        batched=True,
+        remove_columns=dataset["train"].column_names,
+    )
+
+    return tokenized_dataset

src/qa/inference.py

@@ -0,0 +1,206 @@
+import torch
+import faiss
+import json
+import sqlite3
+import re
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+
+
+class LegalQAEngine:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+
+        # ---- Load QA model ----
+        self.tokenizer = AutoTokenizer.from_pretrained("outputs/qa_model/final")
+        self.qa_model = AutoModelForQuestionAnswering.from_pretrained(
+            "outputs/qa_model/final"
+        ).to(self.device)
+        self.qa_model.eval()
+
+        # ---- Load retriever ----
+        self.embedder = SentenceTransformer("BAAI/bge-base-en-v1.5", device=self.device)
+        self.index = faiss.read_index("data/processed/faiss/faiss_index.bin")
+
+        with open("data/processed/embeddings/paragraph_ids.json", encoding="utf-8") as f:
+            self.para_ids = json.load(f)
+
+        self.db = sqlite3.connect("data/processed/indexed/paragraphs.db")
+        self.cursor = self.db.cursor()
+
+        print("✓ Enhanced QA inference system ready")
+
+    # ------------------------------------------------------------------
+    # TEXT NORMALIZATION (critical for PDF artifacts)
+    # ------------------------------------------------------------------
+    def _normalize(self, text: str) -> str:
+        text = text.lower()
+        text = re.sub(r"\s+", " ", text)
+        return text.strip()
+
+    # ------------------------------------------------------------------
+    # REFUTED CLAUSE DETECTION (Article 21 FIX)
+    # ------------------------------------------------------------------
+    def _is_refuted_clause(self, answer_text, paragraph_text):
+        para = self._normalize(paragraph_text)
+        ans = self._normalize(answer_text)
+
+        # Patterns like:
+        # "it is not correct to say, ..., that X"
+        # "it cannot be said, ..., that X"
+        refutation_regexes = [
+            r"not correct to say.*?that\s+(.*?)(?:\.|,)",
+            r"cannot be said.*?that\s+(.*?)(?:\.|,)",
+        ]
+
+        for pattern in refutation_regexes:
+            matches = re.findall(pattern, para)
+            for refuted_prop in matches:
+                # If answer is part of the refuted proposition → block
+                if ans in refuted_prop:
+                    return True
+
+        return False
+
+
+    # ------------------------------------------------------------------
+    # RETRIEVAL
+    # ------------------------------------------------------------------
+    def retrieve_paragraphs(self, question, top_k=8):
+        q_emb = self.embedder.encode(
+            [question], normalize_embeddings=True, convert_to_numpy=True
+        )
+        scores, indices = self.index.search(q_emb, top_k)
+
+        results = []
+        for score, idx in zip(scores[0], indices[0]):
+            para_id = self.para_ids[idx]
+            self.cursor.execute(
+                "SELECT judgment_id, page_no, text FROM paragraphs WHERE id = ?",
+                (para_id,),
+            )
+            row = self.cursor.fetchone()
+            if row:
+                judgment_id, page_no, text = row
+                results.append(
+                    {
+                        "judgment_id": judgment_id,
+                        "page_no": page_no,
+                        "text": text,
+                        "retrieval_score": float(score),
+                    }
+                )
+        return results
+
+    # ------------------------------------------------------------------
+    # ANSWERING
+    # ------------------------------------------------------------------
+    def answer_question(self, question, top_k=8, max_answers=2):
+        paragraphs = self.retrieve_paragraphs(question, top_k)
+        candidates = []
+
+        for para in paragraphs:
+            inputs = self.tokenizer(
+                question,
+                para["text"],
+                return_tensors="pt",
+                truncation=True,
+                max_length=512,
+            ).to(self.device)
+
+            with torch.no_grad():
+                outputs = self.qa_model(**inputs)
+
+            start_logits = outputs.start_logits[0]
+            end_logits = outputs.end_logits[0]
+
+            token_type_ids = inputs["token_type_ids"][0].tolist()
+            question_end = token_type_ids.index(1)
+
+            top_starts = torch.topk(start_logits, k=5).indices
+            top_ends = torch.topk(end_logits, k=5).indices
+
+            for s in top_starts:
+                for e in top_ends:
+                    if e < s or (e - s) > 80:
+                        continue
+
+                    # ❌ Block question echo
+                    if s < question_end:
+                        continue
+
+                    answer_tokens = inputs["input_ids"][0][s : e + 1]
+                    answer_text = self.tokenizer.decode(
+                        answer_tokens, skip_special_tokens=True
+                    ).strip()
+
+                    words = answer_text.split()
+                    if len(words) < 8:
+                        continue
+
+                    # ❌ Block refuted propositions
+                    if self._is_refuted_clause(answer_text, para["text"]):
+                        continue
+
+                    score = start_logits[s].item() + end_logits[e].item()
+
+                    # Doctrinal boost
+                    if any(
+                        k in answer_text.lower()
+                        for k in ["the court", "held that", "it is clear that", "the law"]
+                    ):
+                        score += 1.5
+
+                    candidates.append(
+                        {
+                            "answer": answer_text,
+                            "confidence": score,
+                            "judgment_id": para["judgment_id"],
+                            "page_no": para["page_no"],
+                            "paragraph": para["text"],
+                            "retrieval_score": para["retrieval_score"],
+                        }
+                    )
+
+        # ---- Deduplicate answers ----
+        seen = set()
+        final = []
+        for c in sorted(candidates, key=lambda x: x["confidence"], reverse=True):
+            key = self._normalize(c["answer"])
+            if key not in seen:
+                seen.add(key)
+                final.append(c)
+
+        return final[:max_answers]
+
+
+# ----------------------------------------------------------------------
+# DEMO
+# ----------------------------------------------------------------------
+if __name__ == "__main__":
+    qa = LegalQAEngine()
+
+    questions = [
+        "What is the scope of Article 21?",
+        "What are the conditions for granting anticipatory bail?",
+        "What is the burden of proof in criminal law?",
+    ]
+
+    for q in questions:
+        print("\n" + "=" * 90)
+        print(f"QUESTION: {q}")
+        print("=" * 90)
+
+        answers = qa.answer_question(q)
+
+        for i, ans in enumerate(answers, 1):
+            print(f"\nANSWER {i}:")
+            print(ans["answer"])
+            print(
+                f"\nSOURCE: {ans['judgment_id']} | Page: {ans['page_no']}"
+            )
+            print(f"Retrieval score: {ans['retrieval_score']:.3f}")
+            print(f"Confidence score: {ans['confidence']:.2f}")
+            print("\nPARAGRAPH:")
+            print(ans["paragraph"][:700] + "...")

src/qa/model.py

@@ -0,0 +1,8 @@
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+
+MODEL_NAME = "nlpaueb/legal-bert-base-uncased"
+
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForQuestionAnswering.from_pretrained(MODEL_NAME)
+    return tokenizer, model

src/qa/monitor_training.py

@@ -0,0 +1,21 @@
+"""Monitor training progress"""
+import json
+from pathlib import Path
+import time
+
+log_file = Path("outputs/qa_model/trainer_state.json")
+
+print("Monitoring training progress...\n")
+
+while True:
+    if log_file.exists():
+        with open(log_file) as f:
+            state = json.load(f)
+        
+        if 'log_history' in state:
+            latest = state['log_history'][-1]
+            print(f"Step {latest.get('step', 0):5d} | "
+                  f"Loss: {latest.get('loss', 0):.4f} | "
+                  f"Eval Loss: {latest.get('eval_loss', 'N/A')}")
+    
+    time.sleep(10)

src/qa/train.py

@@ -0,0 +1,42 @@
+from transformers import TrainingArguments, Trainer
+from model import load_model
+from dataset import load_and_prepare_dataset
+
+def main():
+    tokenizer, model = load_model()
+    dataset = load_and_prepare_dataset(tokenizer)
+
+    training_args = TrainingArguments(
+        output_dir="outputs/qa_model",
+        eval_strategy="steps",          # ✅ correct API
+        eval_steps=1000,
+        learning_rate=3e-5,
+        per_device_train_batch_size=8,
+        per_device_eval_batch_size=8,
+        num_train_epochs=2,
+        weight_decay=0.01,
+        fp16=True,
+        logging_steps=500,
+        save_steps=2000,
+        save_total_limit=2,
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["validation"],
+        tokenizer=tokenizer,
+    )
+
+    trainer.train()
+
+    # Save final model + tokenizer
+    trainer.save_model("outputs/qa_model/final")
+    tokenizer.save_pretrained("outputs/qa_model/final")
+
+if __name__ == "__main__":
+    main()

src/rag/query_engine.py

@@ -0,0 +1,326 @@
+"""RAG Query Engine with LLM"""
+
+# pyrefly: ignore [missing-import]
+import faiss
+import json
+import sqlite3
+import re
+# pyrefly: ignore [missing-import]
+from sentence_transformers import SentenceTransformer, CrossEncoder
+import os
+from groq import Groq
+from dotenv import load_dotenv
+
+from src.summarization.ranker import ImportanceRanker
+from src.summarization.utils import split_sentences
+
+load_dotenv()
+
+class QueryEngine:
+    
+    def __init__(self):
+        print("Loading RAG components...")
+        
+        # FAISS + SQLite + Embeddings
+        self.index = faiss.read_index("data/processed/faiss/faiss_index.bin")
+        
+        with open("data/processed/embeddings/paragraph_ids.json") as f:
+            self.para_ids = json.load(f)
+        
+        self.model = SentenceTransformer("BAAI/bge-base-en-v1.5")
+        self.reranker = CrossEncoder("BAAI/bge-reranker-base")
+        self.importance_ranker = ImportanceRanker("outputs/summarization/final")
+        
+    def _get_db(self):
+        return sqlite3.connect("data/processed/indexed/paragraphs.db")
+        
+        # LLM Setup
+        api_key = os.getenv('GROQ_API_KEY')
+        if not api_key:
+            raise ValueError("GROQ_API_KEY not set")
+        
+        self.llm = Groq(api_key=api_key)
+        self.llm_model = 'llama-3.1-8b-instant'  
+        
+        print(f"✓ Ready with {self.index.ntotal:,} vectors")
+        print(f"✓ LLM: Groq (Llama 3.1 8B)")
+    
+    def search(self, query: str, top_k: int = 5):
+        """Hybrid Search: FAISS (Dense) + SQLite FTS5 (BM25) with RRF"""
+        
+        # --- 1. Dense Search (FAISS) ---
+        query_vec = self.model.encode([query], normalize_embeddings=True)
+        dense_scores, dense_indices = self.index.search(query_vec, k=top_k * 2) # Fetch extra for fusion
+        
+        dense_results = []
+        for rank, (score, idx) in enumerate(zip(dense_scores[0], dense_indices[0])):
+            para_id = self.para_ids[idx]
+            dense_results.append({'id': para_id, 'score': float(score), 'rank': rank + 1})
+            
+        # --- 2. Keyword Search (SQLite FTS5 BM25) ---
+        db = self._get_db()
+        cursor = db.cursor()
+        
+        # FTS5 requires a specific syntax. A raw string means AND.
+        # We want BM25 behavior (OR), so we clean punctuation and join words with OR.
+        import re
+        clean_query = re.sub(r'[^\w\s]', '', query)
+        fts_query = " OR ".join(clean_query.split())
+        
+        try:
+            cursor.execute(f"""
+                SELECT id, bm25(paragraphs_fts) as bm25_score
+                FROM paragraphs_fts 
+                WHERE paragraphs_fts MATCH ?
+                ORDER BY bm25_score LIMIT ?
+            """, (fts_query, top_k * 2))
+            fts_rows = cursor.fetchall()
+            
+            keyword_results = []
+            for rank, row in enumerate(fts_rows):
+                keyword_results.append({'id': row[0], 'score': float(row[1]), 'rank': rank + 1})
+        except sqlite3.OperationalError:
+            # Fallback if query syntax is too complex for basic MATCH or FTS table missing
+            keyword_results = []
+
+        # --- 3. Reciprocal Rank Fusion (RRF) ---
+        # RRF Score = 1 / (k + rank)  where k is usually 60
+        k = 60
+        rrf_scores = {}
+        
+        # Add dense scores
+        for res in dense_results:
+            pid = res['id']
+            rrf_scores[pid] = rrf_scores.get(pid, 0.0) + (1.0 / (k + res['rank']))
+            
+        # Add keyword scores
+        for res in keyword_results:
+            pid = res['id']
+            rrf_scores[pid] = rrf_scores.get(pid, 0.0) + (1.0 / (k + res['rank']))
+            
+        # Sort by RRF score descending
+        # Fetch a larger pool of candidates for reranking
+        candidate_pool_size = top_k * 3
+        sorted_rrf = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)[:candidate_pool_size]
+        
+        # --- 4. Fetch Details & Rerank (Cross-Encoder) ---
+        candidates = []
+        for pid, rrf_score in sorted_rrf:
+            cursor.execute(
+                "SELECT judgment_id, text, page_no FROM paragraphs WHERE id = ?",
+                (pid,)
+            )
+            row = cursor.fetchone()
+            
+            if row:
+                candidates.append({
+                    'rrf_score': rrf_score,
+                    'judgment_id': row[0],
+                    'text': row[1],
+                    'page_no': row[2],
+                    'id': pid
+                })
+        
+        if not candidates:
+            db.close()
+            return []
+            
+        db.close()
+
+        # --- 5. Final Rerank (Cross-Encoder) ---
+            
+        # Prepare inputs for cross-encoder: list of [query, document_text]
+        cross_inp = [[query, doc['text']] for doc in candidates]
+        rerank_scores = self.reranker.predict(cross_inp)
+        
+        # Attach scores and sort
+        for i, score in enumerate(rerank_scores):
+            candidates[i]['score'] = float(score)  # Use cross-encoder score as final score
+            
+        candidates = sorted(candidates, key=lambda x: x['score'], reverse=True)
+        
+        return candidates[:top_k]
+    
+    def generate_answer(self, question: str, context: str, sources: list = [], chat_history: list = None):
+        """Generate answer using Groq LLM with strict Legal Guardrails.
+        Sources list is injected into the prompt so the LLM can ONLY cite
+        what was actually retrieved — no hallucinated references.
+        """
+        chat_history = chat_history or []
+        chat_history = chat_history[-6:]  # Cap to last 3 turns
+        # Build a numbered source registry for the LLM
+        source_registry = ""
+        for i, s in enumerate(sources, 1):
+            source_registry += f"[{i}] {s.get('judgment_id', 'Unknown')}\n"
+
+        prompt = f"""You are a strict, brilliant legal research assistant specializing in Indian Supreme Court judgments.
+
+GUARDRAIL: You MUST ONLY answer questions related to law, legal processes, or the provided context.
+If the question is entirely unrelated to law (e.g., "how to bake a cake"), reply EXACTLY with:
+"I am a legal AI assistant. I can only answer questions related to law."
+
+CITATION RULES — THIS IS CRITICAL:
+- You may ONLY cite sources from the APPROVED SOURCE LIST below.
+- Do NOT cite any case from your training memory that is not in the APPROVED SOURCE LIST.
+- If you cite a case not in this list, you are hallucinating and failing your task.
+- Use [1], [2], [3] etc. to refer to sources from the list below.
+
+APPROVED SOURCE LIST (cite ONLY these):
+{source_registry}
+CONTEXT (retrieved paragraphs):
+{context}
+
+QUESTION: {question}
+
+INSTRUCTIONS:
+- Provide a detailed, comprehensive legal answer in a professional conversational tone.
+- Explain concepts clearly so a lawyer finds it extremely useful.
+- Cite ONLY from the APPROVED SOURCE LIST above using [1], [2], [3] format.
+- Use proper legal terminology.
+- Do NOT invent case names, citations, or dates.
+- TEMPORAL AWARENESS: Look at the years in the judgment titles (e.g. 2023_CaseName). Newer judgments (e.g. 2023) supersede older judgments (e.g. 2010). If the retrieved context contains conflicting rulings, you MUST prioritize the newer judgment and explicitly warn the user that the older precedent may have been superseded.
+
+ANSWER:"""
+
+        messages = chat_history.copy()
+        messages.append({"role": "user", "content": prompt})
+
+        response = self.llm.chat.completions.create(
+            model=self.llm_model,
+            messages=messages,
+            temperature=0.2,
+            max_tokens=1024
+        )
+        
+        return response.choices[0].message.content
+    
+    def query(self, question: str, top_k: int = 5, chat_history: list = None):
+        """Main query method"""
+        print(f"\n{'='*70}")
+        print(f"QUERY: {question}")
+        print('='*70)
+        
+        # Search
+        print("\nSearching FAISS index...")
+        results = self.search(question, top_k)
+        
+        print(f"Found {len(results)} relevant paragraphs")
+        
+        # Format context
+        context_parts = []
+        for i, r in enumerate(results, 1):
+            context_parts.append(
+                f"[{i}] {r['judgment_id']}\n{r['text']}"
+            )
+        context = "\n\n".join(context_parts)
+        
+        # Generate answer — pass sources so LLM can only cite what was retrieved
+        print("Generating answer with LLM...")
+        answer = self.generate_answer(question, context, sources=results, chat_history=chat_history)
+        
+        return {
+            'question': question,
+            'answer': answer,
+            'sources': results
+        }
+        
+    def query_with_document(self, question: str, filepath: str, chat_history: list = None):
+        """Queries a specific document. Falls back to global RAG if answer not found."""
+        chat_history = chat_history or []
+        try:
+            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
+                doc_text = f.read()
+        except Exception as e:
+            return {"answer": f"Error reading document: {e}", "sources": []}
+            
+        # JUDGMENTS are usually under 30k chars. Let's take as much as possible.
+        if len(doc_text) > 30000:
+            print("Document exceeds 30k chars. Applying Semantic Truncation...")
+            try:
+                sentences = [s for s in split_sentences(doc_text) if len(s.strip()) > 20]
+                scores = self.importance_ranker.score(sentences)
+                indexed = list(enumerate(zip(sentences, scores)))
+                sorted_by_score = sorted(indexed, key=lambda x: x[1][1], reverse=True)
+                
+                selected_indices = []
+                current_chars = 0
+                for idx, (sentence, score) in sorted_by_score:
+                    if current_chars + len(sentence) > 30000:
+                        continue
+                    selected_indices.append(idx)
+                    current_chars += len(sentence)
+                    if current_chars > 29000:
+                        break
+                
+                # Restore original chronological order
+                top_in_order = sorted([indexed[i] for i in selected_indices], key=lambda x: x[0])
+                doc_text = " ".join(s for _, (s, _) in top_in_order) + "\n\n... [TRUNCATED SEMANTICALLY FOR LLM] ..."
+            except Exception as e:
+                print(f"Semantic Truncation failed: {e}. Falling back to naive truncation.")
+                doc_text = doc_text[:30000] + "\n\n... [TRUNCATED DUE TO SIZE] ..."
+            
+        print(f"--- DOCUMENT QA START ---")
+        print(f"File: {os.path.basename(filepath)}")
+        print(f"Size: {len(doc_text)} chars")
+        print(f"Question: {question}")
+
+        prompt = f"""You are a strict Legal Document Auditor.
+Your ONLY source of information is the text provided below. 
+
+STRICT RULES:
+1. Answer the QUESTION using ONLY the DOCUMENT text.
+2. If the answer is not in the text, say "I cannot find this in the uploaded document."
+3. DO NOT cite external cases (like Venkata Reddy or V.C. Shukla) unless they are explicitly mentioned in the text below.
+4. If you use your own internal knowledge instead of the document, you are failing your task.
+
+DOCUMENT TEXT:
+{doc_text}
+
+QUESTION: {question}
+
+DETAILED ANSWER (citing specific paragraphs if possible):"""
+
+        messages = chat_history.copy()
+        messages.append({"role": "user", "content": prompt})
+
+        response = self.llm.chat.completions.create(
+            model=self.llm_model,
+            messages=messages,
+            temperature=0.1,
+            max_tokens=1024
+        )
+        
+        answer = response.choices[0].message.content.strip()
+            
+        return {
+            'question': question,
+            'answer': answer,
+            'sources': [{'judgment_id': os.path.basename(filepath), 'score': 1.0}]
+        }
+    
+    def close(self):
+        self.db.close()
+
+# Test
+if __name__ == "__main__":
+    engine = QueryEngine()
+    
+    # Test queries
+    queries = [
+        "What are the conditions for granting anticipatory bail?",
+        "Explain the doctrine of legitimate expectation",
+        "What is the burden of proof in criminal cases?"
+    ]
+    
+    for query in queries:
+        response = engine.query(query, top_k=3)
+        
+        print(f"\nANSWER:\n{response['answer']}\n")
+        
+        print("SOURCES:")
+        for i, src in enumerate(response['sources'], 1):
+            print(f"  [{i}] {src['judgment_id']} (score: {src['score']:.3f})")
+        
+        print("\n" + "="*70 + "\n")
+    
+    engine.close()

src/rag/test_retriever.py

@@ -0,0 +1,75 @@
+"""
+Test RAG retrieval system with SQLite
+"""
+
+import faiss
+import json
+import sqlite3
+from sentence_transformers import SentenceTransformer
+
+def test_retrieval():
+    print("="*70)
+    print("Testing RAG Retrieval")
+    print("="*70)
+    
+    # Load FAISS index
+    index = faiss.read_index("data/processed/faiss/faiss_index.bin")
+    
+    # Load paragraph IDs
+    with open("data/processed/embeddings/paragraph_ids.json", encoding="utf-8") as f:
+        para_ids = json.load(f)
+    
+    # Connect to SQLite
+    conn = sqlite3.connect("data/processed/indexed/paragraphs.db")
+    cursor = conn.cursor()
+    
+    # Load embedding model (MUST match indexing)
+    model = SentenceTransformer("BAAI/bge-base-en-v1.5")
+    
+    print(f"✓ Ready with {index.ntotal:,} vectors\n")
+    
+    queries = [
+        "right to privacy under Article 21",
+        "anticipatory bail conditions",
+        "burden of proof in criminal cases",
+        "doctrine of legitimate expectation"
+    ]
+    
+    for query in queries:
+        print(f"\n{'='*70}")
+        print(f"QUERY: {query}")
+        print('='*70)
+        
+        query_vec = model.encode(
+            [query],
+            normalize_embeddings=True,
+            convert_to_numpy=True
+        )
+        
+        scores, indices = index.search(query_vec, k=3)
+        
+        for rank, (score, idx) in enumerate(zip(scores[0], indices[0]), 1):
+            if idx < 0 or idx >= len(para_ids):
+                continue
+            
+            para_id = para_ids[idx]
+            
+            cursor.execute(
+                "SELECT judgment_id, page_no, text FROM paragraphs WHERE id = ?",
+                (para_id,)
+            )
+            row = cursor.fetchone()
+            
+            if not row:
+                continue
+            
+            judgment_id, page_no, text = row
+            
+            print(f"\n[{rank}] Score: {score:.4f}")
+            print(f"Source: {judgment_id} | Page: {page_no}")
+            print(f"Text: {text[:300]}...")
+    
+    conn.close()
+
+if __name__ == "__main__":
+    test_retrieval()

src/segmentation/annotate_paragraphs.py

@@ -0,0 +1,146 @@
+# """
+# Annotate paragraphs with legal sections using JudgmentSegmenter
+# Creates paragraph_index_with_sections.jsonl
+# """
+
+# import json
+# from pathlib import Path
+# from collections import defaultdict
+# from tqdm import tqdm
+
+# from judgement_segmenter import JudgmentSegmenter
+
+
+# INPUT_INDEX = Path("data/processed/indexed/paragraph_index.jsonl")
+# OUTPUT_INDEX = Path("data/processed/indexed/paragraph_index_with_sections.jsonl")
+
+
+# def annotate_paragraphs():
+#     print("=" * 70)
+#     print("NyayLens – Annotating Paragraphs with Sections")
+#     print("=" * 70)
+
+#     # Load paragraphs grouped by judgment
+#     judgments = defaultdict(list)
+
+#     with open(INPUT_INDEX, "r", encoding="utf-8") as f:
+#         for line in f:
+#             p = json.loads(line)
+#             judgments[p["judgment_id"]].append(p)
+
+#     print(f"✓ Loaded {len(judgments):,} judgments")
+
+#     segmenter = JudgmentSegmenter()
+
+#     with open(OUTPUT_INDEX, "w", encoding="utf-8") as writer:
+#         for judgment_id, paras in tqdm(judgments.items(), desc="Annotating"):
+#             # Preserve original order
+#             paras = sorted(paras, key=lambda x: (x["page_no"], x["id"]))
+
+#             texts = [p["text"] for p in paras]
+
+#             sections = segmenter.segment(texts)
+
+#             # Default all to unknown
+#             section_labels = [
+#                 ("unknown", 0.0) for _ in paras
+#             ]
+
+#             # Apply section labels
+#             for sec in sections:
+#                 for i in range(sec.start_para_idx, sec.end_para_idx + 1):
+#                     section_labels[i] = (sec.type, sec.confidence)
+
+#             # Write annotated paragraphs
+#             for p, (sec_type, sec_conf) in zip(paras, section_labels):
+#                 p_out = dict(p)
+#                 p_out["section"] = sec_type
+#                 p_out["section_conf"] = sec_conf
+
+#                 writer.write(json.dumps(p_out, ensure_ascii=False) + "\n")
+
+#     print("\n✓ Annotation complete")
+#     print(f"✓ Output written to: {OUTPUT_INDEX}")
+
+
+# if __name__ == "__main__":
+#     annotate_paragraphs()
+"""
+Annotate paragraphs with legal sections using JudgmentSegmenter
+PRESERVES ORIGINAL IDs AND ORDER
+"""
+
+import json
+from pathlib import Path
+from collections import defaultdict
+from tqdm import tqdm
+
+from judgement_segmenter import JudgmentSegmenter
+
+
+INPUT_INDEX = Path("data/processed/indexed/paragraph_index.jsonl")
+OUTPUT_INDEX = Path("data/processed/indexed/paragraph_index_with_sections.jsonl")
+
+
+def annotate_paragraphs():
+    print("=" * 70)
+    print("NyayLens – Annotating Paragraphs with Sections")
+    print("=" * 70)
+
+    # Load paragraphs IN ORIGINAL ORDER
+    all_paragraphs = []
+    with open(INPUT_INDEX, "r", encoding="utf-8") as f:
+        for line in f:
+            all_paragraphs.append(json.loads(line))
+    
+    print(f"✓ Loaded {len(all_paragraphs):,} paragraphs")
+
+    # Group by judgment (preserve index in group)
+    judgments = defaultdict(list)
+    for idx, p in enumerate(all_paragraphs):
+        judgments[p["judgment_id"]].append((idx, p))  # ← Store original index
+
+    segmenter = JudgmentSegmenter()
+    
+    # Create array to store annotations (preserves original order)
+    annotations = [None] * len(all_paragraphs)
+
+    for judgment_id, indexed_paras in tqdm(judgments.items(), desc="Annotating"):
+        # Extract just the paragraphs
+        indices = [ip[0] for ip in indexed_paras]
+        paras = [ip[1] for ip in indexed_paras]
+        
+        # Get texts
+        texts = [p["text"] for p in paras]
+
+        # Segment
+        sections = segmenter.segment(texts)
+
+        # Default labels
+        section_labels = [("unknown", 0.0) for _ in paras]
+
+        # Apply section labels
+        for sec in sections:
+            for i in range(sec.start_para_idx, sec.end_para_idx + 1):
+                if i < len(section_labels):
+                    section_labels[i] = (sec.type, sec.confidence)
+
+        # Store annotations in ORIGINAL positions
+        for orig_idx, p, (sec_type, sec_conf) in zip(indices, paras, section_labels):
+            p_out = dict(p)  # Copy original
+            p_out["section"] = sec_type
+            p_out["section_conf"] = sec_conf
+            annotations[orig_idx] = p_out
+
+    # Write in ORIGINAL order
+    print("\nWriting annotated paragraphs...")
+    with open(OUTPUT_INDEX, "w", encoding="utf-8") as writer:
+        for p_out in annotations:
+            writer.write(json.dumps(p_out, ensure_ascii=False) + "\n")
+
+    print(f"✓ Output written to: {OUTPUT_INDEX}")
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    annotate_paragraphs()

src/segmentation/check.py

@@ -0,0 +1,41 @@
+"""Test section-aware retrieval"""
+import faiss
+import json
+import sqlite3
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+# Load
+index = faiss.read_index("data/processed/faiss/faiss_index.bin")
+with open("data/processed/embeddings/paragraph_ids.json") as f:
+    para_ids = json.load(f)
+
+db = sqlite3.connect("data/processed/indexed/paragraphs.db")
+cursor = db.cursor()
+
+model = SentenceTransformer("BAAI/bge-base-en-v1.5")
+
+# Test query
+query = "What were the facts of the case?"
+query_vec = model.encode([query], normalize_embeddings=True)
+
+# Search
+scores, indices = index.search(query_vec, k=10)
+
+print(f"Query: {query}\n")
+print("Top results with sections:")
+
+for i, (score, idx) in enumerate(zip(scores[0], indices[0]), 1):
+    para_id = para_ids[idx]
+    
+    cursor.execute(
+        "SELECT judgment_id, section, section_conf, text FROM paragraphs WHERE id = ?",
+        (para_id,)
+    )
+    row = cursor.fetchone()
+    
+    print(f"\n[{i}] Score: {score:.3f} | Section: {row[1]} (conf={row[2]:.2f})")
+    print(f"    Case: {row[0]}")
+    print(f"    Text: {row[3][:100]}...")
+
+db.close()

src/segmentation/judgement_segmenter.py

@@ -0,0 +1,195 @@
+"""
+Enhanced Judgment Segmenter (FIXED)
+Segments judgments into: Facts, Issues, Arguments, Analysis, Decision
+"""
+
+import re
+import os
+import logging
+from typing import List, Dict, Tuple
+from dataclasses import dataclass
+
+try:
+    from transformers import pipeline
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Section:
+    type: str                  # facts/issues/arguments/analysis/decision/unknown
+    text: str
+    start_para_idx: int
+    end_para_idx: int
+    confidence: float
+
+
+class JudgmentSegmenter:
+
+    MARKERS = {
+        'facts': [
+            r'\bbrief\s+facts?\b',
+            r'\bfactual\s+(matrix|background)\b',
+            r'\bcircumstances\s+of\s+the\s+case\b',
+            r'\bbackground\b',
+        ],
+        'issues': [
+            r'\bissues?\s+(for|of)\s+(consideration|determination)\b',
+            r'\bsubstantial\s+questions?\b',
+            r'\bpoints?\s+for\s+consideration\b',
+            r'\bquestions?\s+framed\b',
+        ],
+        'arguments': [
+            r'\blearned\s+counsel\b',
+            r'\bsubmissions?\b',
+            r'\b(argued|submitted|contended)\b',
+            r'\bon\s+behalf\s+of\b',
+        ],
+        'analysis': [
+            r'\bwe\s+have\s+(considered|examined|analysed)\b',
+            r'\bthe\s+court\s+(finds|observes|notes|holds)\b',
+            r'\bin\s+our\s+(view|opinion)\b',
+            r'\bit\s+is\s+clear\s+that\b',
+        ],
+        'decision': [
+            r'\b(appeal|petition|writ)\s+is\s+(allowed|dismissed)\b',
+            r'\baccordingly\b',
+            r'\bwe\s+direct\b',
+            r'\bheld\s*:\b',
+            r'\border\b',
+        ]
+    }
+
+    def __init__(self, model_path: str = "models/segmentation_model"):
+        """Initialize segmenter, preferring ML model if available, else Regex fallback"""
+        self.use_ml = False
+        self.classifier = None
+        
+        if TRANSFORMERS_AVAILABLE and os.path.exists(model_path):
+            try:
+                logger.info(f"Loading ML Segmentation model from {model_path}...")
+                self.classifier = pipeline("text-classification", model=model_path, device=-1)
+                self.use_ml = True
+                logger.info("✓ ML Segmenter loaded successfully.")
+            except Exception as e:
+                logger.warning(f"Failed to load ML model, falling back to Regex: {e}")
+        else:
+            logger.info("ML model not found or transformers not installed. Using Regex fallback.")
+
+    def detect_section(self, para: str, position_ratio: float) -> Tuple[str, float]:
+        """
+        Detect section type for a paragraph
+        Returns: (section_type, confidence)
+        """
+        para_lower = para.lower()
+        best_type = 'unknown'
+        best_conf = 0.0
+
+        for sec_type, patterns in self.MARKERS.items():
+            for pattern in patterns:
+                if re.search(pattern, para_lower):
+                    conf = 0.6
+
+                    # Position-based bias
+                    if sec_type == 'facts' and position_ratio < 0.30:
+                        conf += 0.2
+                    elif sec_type == 'decision' and position_ratio > 0.70:
+                        conf += 0.3
+
+                    # Strong anchor near paragraph start
+                    if re.search(pattern, para_lower[:120]):
+                        conf += 0.2
+
+                    conf = min(conf, 1.0)
+
+                    if conf > best_conf:
+                        best_type = sec_type
+                        best_conf = conf
+
+        return best_type, best_conf
+        
+    def detect_section_ml(self, para: str) -> Tuple[str, float]:
+        """Detect using HuggingFace classifier"""
+        if not para.strip() or not self.classifier:
+            return "unknown", 0.0
+            
+        # Truncate to max length to avoid tokenization errors
+        truncated = para[:512]
+        result = self.classifier(truncated)[0]
+        
+        # Assume labels are like LABEL_FACTS, LABEL_ISSUES or directly facts, issues
+        label = result['label'].lower().replace('label_', '')
+        score = result['score']
+        
+        # Enforce confidence threshold
+        if score < 0.5:
+            return "unknown", score
+            
+        return label, score
+
+    def segment(self, paragraph_texts: List[str]) -> List[Section]:
+        """
+        Segment judgment based on paragraph list (INDEX-ALIGNED)
+        """
+        if not paragraph_texts:
+            return []
+
+        sections: List[Section] = []
+
+        current_type = 'unknown'
+        current_paras = []
+        current_conf = 0.0
+        start_idx = 0
+
+        total = len(paragraph_texts)
+
+        for i, para in enumerate(paragraph_texts):
+            position_ratio = i / max(total, 1)
+            
+            if self.use_ml:
+                sec_type, conf = self.detect_section_ml(para)
+            else:
+                sec_type, conf = self.detect_section(para, position_ratio)
+
+            # Fallback: early unknown paragraphs are likely facts
+            if sec_type == 'unknown' and position_ratio < 0.30 and i > 0:
+                sec_type = 'facts'
+                conf = 0.4
+
+            # Section boundary
+            if conf > 0.4 and sec_type != current_type:
+                if current_paras:
+                    sections.append(
+                        Section(
+                            type=current_type,
+                            text="\n\n".join(current_paras),
+                            start_para_idx=start_idx,
+                            end_para_idx=i - 1,
+                            confidence=round(current_conf, 2)
+                        )
+                    )
+
+                current_type = sec_type
+                current_paras = [para]
+                current_conf = conf
+                start_idx = i
+            else:
+                current_paras.append(para)
+                current_conf = max(current_conf, conf)
+
+        # Final section
+        if current_paras:
+            sections.append(
+                Section(
+                    type=current_type,
+                    text="\n\n".join(current_paras),
+                    start_para_idx=start_idx,
+                    end_para_idx=total - 1,
+                    confidence=round(current_conf, 2)
+                )
+            )
+
+        return sections

src/summarization/composer.py

@@ -0,0 +1,17 @@
+# src/summarization/composer.py
+def compose(sentences, scores, top_k=5):
+    """
+    Select top-k sentences by Legal-BERT score, then RESTORE their original
+    document order before returning. This gives PEGASUS a coherent narrative
+    instead of a randomly ordered bag of sentences.
+    """
+    # Tag each sentence with its original index
+    indexed = list(enumerate(zip(sentences, scores)))
+    
+    # Pick top-k by score
+    top = sorted(indexed, key=lambda x: x[1][1], reverse=True)[:top_k]
+    
+    # Re-sort by original document position for narrative coherence
+    top_in_order = sorted(top, key=lambda x: x[0])
+    
+    return [s for _, (s, _) in top_in_order]

src/summarization/dataset.py

@@ -0,0 +1,74 @@
+# src/summarization/dataset.py
+import re
+from pathlib import Path
+from datasets import Dataset
+from transformers import AutoTokenizer
+from tqdm import tqdm
+
+IMPORTANT_PATTERNS = [
+    r"\bheld\b",
+    r"\bwe conclude\b",
+    r"\btherefore\b",
+    r"\bappeal is (allowed|dismissed)\b",
+    r"\bsubstantial question\b",
+    r"\baccordingly\b",
+]
+
+def sentence_split(text):
+    return re.split(r'(?<=[.!?])\s+', text)
+
+def is_important(sentence):
+    s = sentence.lower()
+    return any(re.search(p, s) for p in IMPORTANT_PATTERNS)
+
+def build_dataset(text_dir, tokenizer_name, limit=None):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    samples = []
+
+    files = list(Path(text_dir).glob("*.txt"))
+    if limit:
+        files = files[:limit]
+        
+    for file in tqdm(files, desc="Processing judgments"):
+        judgment_id = file.stem
+        text = file.read_text(encoding="utf-8", errors="ignore")
+
+        sentences = sentence_split(text)
+        for sent in sentences:
+            sent = sent.strip()
+            if len(sent) < 40:
+                continue
+
+            samples.append({
+                "text": sent,
+                "label": int(is_important(sent)),
+                "judgment_id": judgment_id
+            })
+
+    dataset = Dataset.from_list(samples)
+
+    def tokenize(batch):
+        return tokenizer(
+            batch["text"],
+            truncation=True,
+            padding="max_length",
+            max_length=256
+        )
+
+    return dataset.map(tokenize, batched=True)
+
+if __name__ == "__main__":
+    print("Started dataset building...")
+    # Using a limit of 1000 for training, can be increased later
+    # 1000 judgments will yield ~50k-100k sentences, good for fine-tuning
+    ds = build_dataset(
+        "data/processed/extracted/texts",
+        "nlpaueb/legal-bert-base-uncased",
+        limit=1000
+    )
+    print("Tokenizing dataset... this may take a moment.")
+    print(f"Total sentences extracted: {len(ds)}")
+    
+    print("Saving to Disk...")
+    ds.save_to_disk("data/processed/summarization_dataset")
+    print("✓ Dataset ready at data/processed/summarization_dataset")

src/summarization/inference.py

@@ -0,0 +1,129 @@
+# src/summarization/inference.py
+import sys
+import os
+import re
+from pathlib import Path
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
+
+from src.summarization.ranker import ImportanceRanker
+from src.summarization.utils import split_sentences
+from src.segmentation.judgement_segmenter import JudgmentSegmenter
+from transformers import PegasusTokenizer, PegasusForConditionalGeneration
+import torch
+
+# ── Model ──────────────────────────────────────────────────────────────────────
+MODEL_NAME = "nsi319/legal-pegasus"
+print(f"\nLoading Abstractive Model ({MODEL_NAME})...")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pegasus_tokenizer = PegasusTokenizer.from_pretrained(MODEL_NAME)
+pegasus_model = PegasusForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
+print(f"✓ Legal-PEGASUS loaded on {device.upper()}")
+
+
+def _pegasus_generate(text: str, max_length: int = 300, min_length: int = 100) -> str:
+    """Run Legal-PEGASUS on a block of text and return the decoded summary."""
+    inputs = pegasus_tokenizer(
+        [text],
+        max_length=1024,
+        truncation=True,
+        padding=True,
+        return_tensors="pt"
+    ).to(device)
+
+    outputs = pegasus_model.generate(
+        inputs["input_ids"],
+        max_length=max_length,
+        min_length=min_length,
+        num_beams=4,                # Reduced from 8 for 2x speedup on CPU
+        length_penalty=1.2,
+        no_repeat_ngram_size=3,
+        repetition_penalty=1.3,
+        early_stopping=True,
+    )
+    decoded = pegasus_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return re.sub(r'\s+', ' ', decoded.replace("<n>", " ")).strip()
+
+
+def summarize(judgment_file: str) -> dict:
+    """
+    Speed-Optimized Two-Pass Pipeline:
+    1. Case Overview: Legal-BERT (Extraction) -> Legal-PEGASUS (Abstraction) [1 Pass]
+    2. Detailed Sections: Legal-BERT (Extraction) -> Direct Output [No Abstraction pass to save 5+ minutes]
+    """
+    text = Path(judgment_file).read_text(encoding="utf-8", errors="ignore")
+
+    # ── Step 1: Global sentence extraction ─────────────────────────────────────
+    all_sentences = [s for s in split_sentences(text) if len(s.strip()) > 40]
+    if not all_sentences:
+        return {"overview": "Could not extract readable text."}
+
+    ranker = ImportanceRanker("outputs/summarization/final")
+    scores = ranker.score(all_sentences)
+
+    # Token-Aware Global Overview Extract (Limit to ~950 tokens for Pegasus)
+    indexed = list(enumerate(zip(all_sentences, scores)))
+    sorted_by_score = sorted(indexed, key=lambda x: x[1][1], reverse=True)
+    
+    selected_indices = []
+    current_tokens = 0
+    MAX_TOKENS = 950
+    
+    for idx, (sentence, score) in sorted_by_score:
+        tokens = len(pegasus_tokenizer.encode(sentence, add_special_tokens=False))
+        if current_tokens + tokens > MAX_TOKENS:
+            continue
+            
+        selected_indices.append(idx)
+        current_tokens += tokens
+        if current_tokens >= MAX_TOKENS - 20:
+            break
+            
+    # Restore chronological order
+    top_in_order = sorted([indexed[i] for i in selected_indices], key=lambda x: x[0])
+    global_extract = " ".join(s for _, (s, _) in top_in_order)
+
+    # ── Pass 1: Abstractive Overview (The only heavy pass) ────────────────────
+    print("Generating Case Overview (Abstractive)...")
+    overview = _pegasus_generate(global_extract, max_length=250, min_length=80)
+
+    # ── Pass 2: Extractive Section Breakdown (Instant) ────────────────────────
+    segmenter = JudgmentSegmenter()
+    paragraphs = [p.strip() for p in text.split("\n\n") if len(p.strip()) > 20]
+    sections = segmenter.segment(paragraphs)
+
+    final_summary = {"overview": overview}
+
+    print("Generating Section Breakdowns (Extractive - Instant)...")
+    for section in sections:
+        sec_type = section.type.lower()
+        if sec_type == 'unknown': continue
+
+        sentences = [s for s in split_sentences(section.text) if len(s.strip()) > 40]
+        if not sentences: continue
+
+        sec_scores = ranker.score(sentences)
+        # Select top 3 per section for readability
+        s_indexed = list(enumerate(zip(sentences, sec_scores)))
+        top_k = sorted(s_indexed, key=lambda x: x[1][1], reverse=True)[:3]
+        top_k_ordered = sorted(top_k, key=lambda x: x[0])
+        
+        # We use original sentences here to save 5-10 minutes of CPU time
+        final_summary[sec_type] = " ".join(s for _, (s, _) in top_k_ordered)
+
+    return final_summary
+
+
+if __name__ == "__main__":
+    file = list(Path("data/processed/extracted/texts").glob("*.txt"))[0]
+    print(f"\nProcessing {file.name}...")
+    result = summarize(file)
+
+    print("\n\nCOMPREHENSIVE LEGAL SUMMARY (Global Legal-BERT + Legal-PEGASUS)\n" + "=" * 80)
+    print("\n[CASE OVERVIEW]")
+    print(result.get("overview", ""))
+    for sec in ['facts', 'issues', 'arguments', 'analysis', 'decision']:
+        if sec in result:
+            print(f"\n[{sec.upper()}]")
+            print(result[sec])
+    print("\n" + "=" * 80)

src/summarization/model.py

@@ -0,0 +1,25 @@
+# src/summarization/model.py
+from transformers import AutoModel
+import torch
+import torch.nn as nn
+
+class SentenceRanker(nn.Module):
+    def __init__(self, model_name):
+        super().__init__()
+        self.encoder = AutoModel.from_pretrained(model_name)
+        self.classifier = nn.Linear(self.encoder.config.hidden_size, 1)
+
+    def forward(self, input_ids, attention_mask, labels=None, **kwargs):
+        out = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **kwargs
+        )
+        cls = out.last_hidden_state[:, 0]
+        logits = self.classifier(cls).squeeze(-1)
+
+        loss = None
+        if labels is not None:
+            loss = nn.BCEWithLogitsLoss()(logits, labels.float())
+
+        return {"loss": loss, "logits": logits}

src/summarization/ranker.py

@@ -0,0 +1,40 @@
+# src/summarization/ranker.py
+import torch
+from transformers import AutoTokenizer
+from src.summarization.model import SentenceRanker
+
+class ImportanceRanker:
+    def __init__(self, model_dir, base_model="nlpaueb/legal-bert-base-uncased"):
+        # Load the tokenizer from the base model
+        self.tokenizer = AutoTokenizer.from_pretrained(base_model)
+        
+        # Initialize the custom architecture with base model
+        self.model = SentenceRanker(base_model)
+        
+        # Load fine-tuned weights
+        import os
+        from safetensors.torch import load_file
+        
+        weights_path = os.path.join(model_dir, "model.safetensors")
+        if os.path.exists(weights_path):
+            state_dict = load_file(weights_path)
+            self.model.load_state_dict(state_dict)
+        else:
+            print(f"Warning: Could not find {weights_path}")
+            
+        self.model.eval()
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+
+    def score(self, sentences):
+        inputs = self.tokenizer(
+            sentences,
+            truncation=True,
+            padding=True,
+            return_tensors="pt"
+        ).to(self.device)
+
+        with torch.no_grad():
+            logits = self.model(**inputs)["logits"]
+
+        return logits.sigmoid().cpu().tolist()

src/summarization/train.py

@@ -0,0 +1,41 @@
+# src/summarization/train.py
+
+from datasets import load_from_disk
+from transformers import Trainer, TrainingArguments
+from model import SentenceRanker
+import torch
+
+MODEL_NAME = "nlpaueb/legal-bert-base-uncased"
+
+def main():
+    # Load dataset
+    dataset = load_from_disk("data/processed/summarization_dataset")
+    dataset = dataset.train_test_split(test_size=0.1)
+
+    model = SentenceRanker(MODEL_NAME)
+
+    training_args = TrainingArguments(
+        output_dir="outputs/summarization",
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        num_train_epochs=2,
+        learning_rate=2e-5,
+        logging_steps=500,
+        save_steps=2000,
+        save_total_limit=2,
+        report_to="none",
+        fp16=torch.cuda.is_available()
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["test"]
+    )
+
+    trainer.train()
+    trainer.save_model("outputs/summarization/final")
+
+if __name__ == "__main__":
+    main()

src/summarization/utils.py

@@ -0,0 +1,5 @@
+# src/summarization/utils.py
+import re
+
+def split_sentences(text):
+    return re.split(r'(?<=[.!?])\s+', text)