Add RAG module: PostgreSQL + pgvector in same container

- PostgreSQL 16 + pgvector installed in Docker image
- start.sh initializes DB and starts FastAPI on port 7860
- Full RAG pipeline: TF-IDF/SVD embeddings + pgvector search
- NVIDIA NIM justification (requires NVIDIA_API_KEY secret)
- Multi-model inference: logistic, lgbm, xgb
- train.json (16MB) tracked via git-lfs
- Angular SPA served from FastAPI

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+data/raw/*.json filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -1,5 +1,17 @@
 FROM python:3.11-slim
 
+RUN apt-get update && apt-get install -y \
+    postgresql \
+    postgresql-contrib \
+    libpq-dev \
+    build-essential \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git /pgvector \
+    && cd /pgvector && make && make install \
+    && rm -rf /pgvector
+
 WORKDIR /app
 
 COPY requirements.txt .
@@ -7,8 +19,11 @@ RUN pip install --no-cache-dir -r requirements.txt
 
 COPY api/ api/
 COPY models/ models/
+COPY data/ data/
 COPY frontend/dist/frontend/browser/ frontend/dist/frontend/browser/
+COPY start.sh /start.sh
+RUN chmod +x /start.sh
 
 EXPOSE 7860
 
-CMD ["uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["/start.sh"]

api/inference.py

@@ -1,22 +1,104 @@
 import joblib
+import os
+from fastapi import HTTPException
 
-model = joblib.load(
-    "models/best_logistic.pkl"
-)
+MODELS_DIR = os.path.join(os.path.dirname(__file__), '..', 'models')
 
-vectorizer = joblib.load(
-    "models/tfidf_vectorizer.pkl"
-)
+_cache: dict = {}
 
 
-def predict_claim(text):
+def _path(filename: str) -> str:
+    return os.path.join(MODELS_DIR, filename)
 
-    vector = vectorizer.transform(
-        [text]
-    )
 
-    prediction = model.predict(
-        vector
-    )
+def _available(filename: str) -> bool:
+    return os.path.exists(_path(filename))
 
-    return str(prediction[0])
+
+def _load(filename: str, key: str) -> bool:
+    if not _available(filename):
+        return False
+    try:
+        _cache[key] = joblib.load(_path(filename))
+        return True
+    except Exception as e:
+        print(f"WARNING: could not load {filename}: {e}")
+        return False
+
+
+def get_available_models() -> list[str]:
+    available = []
+    if 'logistic' in _cache: available.append('logistic')
+    if 'lgbm'     in _cache: available.append('lgbm')
+    if 'xgb'      in _cache: available.append('xgb')
+    return available
+
+
+def load_models() -> None:
+    if 'vectorizer' not in _cache:
+        if not _available('tfidf_vectorizer.pkl'):
+            raise RuntimeError(
+                "tfidf_vectorizer.pkl not found in models/. "
+                "Run scripts/train_all_models.py first."
+            )
+        _cache['vectorizer'] = joblib.load(_path('tfidf_vectorizer.pkl'))
+        print("Vectorizer loaded.")
+
+    if 'logistic' not in _cache:
+        if _load('best_logistic.pkl', 'logistic'):
+            print("Logistic model loaded.")
+
+    if 'lgbm' not in _cache:
+        if _load('best_lgbm.pkl', 'lgbm'):
+            print("LightGBM model loaded.")
+
+    if 'xgb' not in _cache:
+        ok1 = _load('best_xgb.pkl',    'xgb')
+        ok2 = _load('xgb_encoder.pkl', 'xgb_encoder')
+        if ok1 and ok2:
+            print("XGBoost model loaded.")
+        elif ok1:
+            del _cache['xgb']
+
+    available = get_available_models()
+    if not available:
+        raise RuntimeError(
+            "No models loaded. Run scripts/train_all_models.py first."
+        )
+    print(f"Models ready: {available}")
+
+
+def predict_claim(text: str, model_id: str = 'logistic') -> dict:
+    if 'vectorizer' not in _cache:
+        raise HTTPException(status_code=503, detail="Models not loaded.")
+
+    vectorizer = _cache['vectorizer']
+    vector = vectorizer.transform([text])
+
+    if model_id == 'lgbm':
+        if 'lgbm' not in _cache:
+            raise HTTPException(status_code=503, detail="LightGBM model not available.")
+        model = _cache['lgbm']
+        prediction   = str(model.predict(vector)[0])
+        proba_values = model.predict_proba(vector)[0]
+        probabilities = {str(c): float(p) for c, p in zip(model.classes_, proba_values)}
+
+    elif model_id == 'xgb':
+        if 'xgb' not in _cache:
+            raise HTTPException(status_code=503, detail="XGBoost model not available.")
+        model   = _cache['xgb']
+        encoder = _cache['xgb_encoder']
+        pred_enc      = model.predict(vector)[0]
+        prediction    = str(encoder.inverse_transform([pred_enc])[0])
+        proba_values  = model.predict_proba(vector)[0]
+        probabilities = {str(c): float(p) for c, p in zip(encoder.classes_, proba_values)}
+
+    else:  # logistic (default)
+        if 'logistic' not in _cache:
+            raise HTTPException(status_code=503, detail="Logistic model not available.")
+        model = _cache['logistic']
+        prediction   = str(model.predict(vector)[0])
+        proba_values = model.predict_proba(vector)[0]
+        probabilities = {str(c): float(p) for c, p in zip(model.classes_, proba_values)}
+
+    return {"prediction": prediction, "probabilities": probabilities}

api/main.py

@@ -1,29 +1,80 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
-from api.inference import predict_claim
-from api.schema import ClaimRequest, PredictionResponse
+from fastapi.middleware.cors import CORSMiddleware
+from api.inference import predict_claim, load_models, get_available_models
+from api.schema import ClaimRequest, PredictionResponse, EvidenceItem
+from api.rag import load_rag_data, get_rag_result
+from dotenv import load_dotenv
+import asyncio
 import os
+from contextlib import asynccontextmanager
 
-app = FastAPI()
+load_dotenv()
+
+STATIC_DIR = "frontend/dist/frontend/browser"
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    print("Loading models...")
+    load_models()
+    print("Building RAG index...")
+    load_rag_data()
+    print("Startup complete.")
+    yield
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 
 @app.get("/health")
 def health():
-    return {"status": "ok"}
+    return {"status": "ok", "available_models": get_available_models()}
+
+
+@app.get("/models")
+def list_models():
+    return {"available_models": get_available_models()}
 
 
 @app.post("/predict")
-def predict_claim_endpoint(claim_request: ClaimRequest):
-    prediction = predict_claim(claim_request.claim)
-    return PredictionResponse(prediction=prediction)
+async def predict_claim_endpoint(claim_request: ClaimRequest):
+    try:
+        claim = claim_request.claim.strip()
+        if not claim:
+            raise HTTPException(status_code=400, detail="Claim text cannot be empty")
 
+        result = await asyncio.to_thread(predict_claim, claim, claim_request.model)
+        predicted_label = result["prediction"]
 
-STATIC_DIR = "frontend/dist/frontend/browser"
+        evidence_raw, justification = await asyncio.to_thread(
+            get_rag_result, claim, predicted_label
+        )
 
-if os.path.exists(STATIC_DIR):
-    app.mount("/assets", StaticFiles(directory=STATIC_DIR), name="assets")
+        evidence = [EvidenceItem(**e) for e in evidence_raw] if evidence_raw else None
 
+        return PredictionResponse(
+            prediction=predicted_label,
+            probabilities=result.get("probabilities"),
+            evidence=evidence,
+            justification=justification,
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing claim: {str(e)}")
+
+
+if os.path.exists(STATIC_DIR):
     @app.get("/{full_path:path}")
     async def serve_spa(full_path: str):
         file_path = os.path.join(STATIC_DIR, full_path)

api/rag.py

@@ -0,0 +1,230 @@
+import os
+import json
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.decomposition import TruncatedSVD
+from sklearn.pipeline import Pipeline
+
+DATA_PATH = os.path.join(os.path.dirname(__file__), '..', 'data', 'raw', 'train.json')
+RAG_DIMS  = 512
+
+_pipeline: Pipeline | None = None
+_pgvector_ready = False
+
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+def _database_url() -> str:
+    return os.environ.get(
+        'DATABASE_URL',
+        'postgresql://mlflow:mlflow123@localhost:5432/mlflow_db'
+    )
+
+
+def _connect():
+    import psycopg2
+    from pgvector.psycopg2 import register_vector
+    conn = psycopg2.connect(_database_url())
+    register_vector(conn)
+    return conn
+
+
+def _fit_pipeline(texts: list[str]) -> Pipeline:
+    pipe = Pipeline([
+        ('tfidf', TfidfVectorizer(max_features=5000, stop_words='english', ngram_range=(1, 2))),
+        ('svd',   TruncatedSVD(n_components=RAG_DIMS, random_state=42)),
+    ])
+    pipe.fit(texts)
+    return pipe
+
+
+def _embed(pipe: Pipeline, texts: list[str]) -> np.ndarray:
+    vecs = pipe.transform(texts)
+    norms = np.linalg.norm(vecs, axis=1, keepdims=True)
+    norms[norms == 0] = 1.0
+    return (vecs / norms).astype(np.float32)
+
+
+# ── startup ───────────────────────────────────────────────────────────────────
+
+def load_rag_data(_main_vectorizer=None) -> None:
+    global _pipeline, _pgvector_ready
+
+    if not os.path.exists(DATA_PATH):
+        print(f"RAG: {DATA_PATH} not found — retrieval disabled.")
+        return
+
+    try:
+        with open(DATA_PATH, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+    except Exception as e:
+        print(f"RAG: failed to load data — {e}")
+        return
+
+    statements = [str(item.get('statement', '')) for item in data]
+    labels     = [str(item.get('label',     '')) for item in data]
+    reasons    = [str(item.get('reason',    '')) for item in data]
+
+    # Always (re)fit the pipeline — fast enough on 10K rows
+    print("RAG: fitting LSA pipeline (TF-IDF 5 000 → SVD 512)…")
+    _pipeline = _fit_pipeline(statements)
+
+    db_url = _database_url()
+    print(f"RAG: connecting to {db_url.split('@')[-1]}…")  # hide credentials
+
+    try:
+        conn = _connect()
+        cur  = conn.cursor()
+
+        cur.execute("SELECT COUNT(*) FROM rag_claims;")
+        count = cur.fetchone()[0]
+
+        if count == 0:
+            print(f"RAG: inserting {len(statements)} claims into pgvector…")
+            vectors = _embed(_pipeline, statements)
+
+            from psycopg2.extras import execute_values
+            rows = [
+                (statements[i], labels[i], reasons[i][:500], vectors[i])
+                for i in range(len(statements))
+            ]
+            execute_values(
+                cur,
+                "INSERT INTO rag_claims (statement, label, reason, embedding) VALUES %s",
+                rows,
+                template="(%s, %s, %s, %s)",
+            )
+            conn.commit()
+            print(f"RAG: {len(rows)} claims indexed in pgvector ✓")
+        else:
+            print(f"RAG: {count} claims already in pgvector — skipping insert.")
+
+        cur.close()
+        conn.close()
+        _pgvector_ready = True
+        print("RAG: ready ✓")
+
+    except Exception as e:
+        print(f"RAG ERROR: {e}")
+        print("RAG: retrieval disabled — run 'docker compose down -v && docker compose up --build' to reset.")
+        _pgvector_ready = False
+
+
+# ── retrieval ─────────────────────────────────────────────────────────────────
+
+def retrieve_similar(claim: str, predicted_label: str, top_k: int = 3) -> list:
+    if not _pgvector_ready or _pipeline is None:
+        return []
+
+    try:
+        query_vec = _embed(_pipeline, [claim])[0]
+
+        conn = _connect()
+        cur  = conn.cursor()
+        cur.execute(
+            """
+            SELECT statement, label, reason,
+                   1 - (embedding <=> %s) AS similarity
+            FROM   rag_claims
+            WHERE  label = %s
+            ORDER  BY embedding <=> %s
+            LIMIT  %s
+            """,
+            (query_vec, predicted_label, query_vec, top_k),
+        )
+        rows = cur.fetchall()
+        cur.close()
+        conn.close()
+
+        return [
+            {
+                'statement':  row[0],
+                'label':      row[1],
+                'reason':     row[2],
+                'similarity': float(row[3]),
+            }
+            for row in rows
+        ]
+    except Exception as e:
+        print(f"RAG: retrieval error — {e}")
+        return []
+
+
+# ── generation (NVIDIA NIM) ───────────────────────────────────────────────────
+
+_NIM_BASE_URL = "https://integrate.api.nvidia.com/v1"
+_NIM_MODEL    = "meta/llama-3.1-8b-instruct"
+
+
+def _generate_justification(claim: str, predicted_label: str, similar: list) -> str | None:
+    if not similar:
+        return None
+
+    nim_key = os.environ.get('NVIDIA_API_KEY')
+    if not nim_key:
+        print("RAG: NVIDIA_API_KEY not set — justification disabled.")
+        return None
+
+    try:
+        from openai import OpenAI
+        client = OpenAI(base_url=_NIM_BASE_URL, api_key=nim_key)
+    except ImportError:
+        return None
+
+    label_fr = {
+        'true':        'VRAI',
+        'mostly-true': 'MAJORITAIREMENT VRAI',
+        'half-true':   'PARTIELLEMENT VRAI',
+        'barely-true': 'À PEINE VRAI',
+        'false':       'FAUX',
+        'pants-fire':  'TOTALEMENT FAUX',
+    }.get(predicted_label, predicted_label.upper())
+
+    examples = []
+    for i, s in enumerate(similar[:3], 1):
+        examples.append(
+            f"Exemple {i} (similarité={s['similarity']:.2f}) :\n"
+            f"  Déclaration : {s['statement']}\n"
+            f"  Raison : {(s['reason'] or '')[:250]}"
+        )
+
+    prompt = (
+        f"Tu es un assistant de vérification des faits.\n"
+        f"Déclaration : « {claim} »\n"
+        f"Classification : {label_fr} ({predicted_label})\n\n"
+        f"Exemples similaires (même classification) dans la base LIAR :\n\n"
+        + '\n\n'.join(examples) +
+        f"\n\nEn 2-3 phrases concises en français, justifie pourquoi cette déclaration est"
+        f" classée « {predicted_label} » en t'appuyant sur les similitudes avec ces exemples."
+        " Réponds directement sans titre ni introduction."
+    )
+
+    try:
+        resp = client.chat.completions.create(
+            model=_NIM_MODEL,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=350,
+            temperature=0.3,
+        )
+        return resp.choices[0].message.content.strip()
+    except Exception as exc:
+        print(f"RAG NIM generation error: {exc}")
+        return None
+
+
+# ── public API ────────────────────────────────────────────────────────────────
+
+def get_rag_result(claim: str, predicted_label: str, _vectorizer=None) -> tuple[list, str | None]:
+    similar       = retrieve_similar(claim, predicted_label)
+    justification = _generate_justification(claim, predicted_label, similar)
+
+    evidence = [
+        {
+            'statement':  s['statement'],
+            'label':      s['label'],
+            'similarity': round(s['similarity'], 3),
+            'reason':     (s['reason'] or '')[:300] or None,
+        }
+        for s in similar
+    ]
+    return evidence, justification

api/schema.py

@@ -1,9 +1,21 @@
 from pydantic import BaseModel
+from typing import Dict, List, Optional
 
 
 class ClaimRequest(BaseModel):
     claim: str
+    model: str = 'logistic'
+
+
+class EvidenceItem(BaseModel):
+    statement: str
+    label: str
+    similarity: float
+    reason: Optional[str] = None
 
 
 class PredictionResponse(BaseModel):
-    prediction: str
+    prediction: str
+    probabilities: Optional[Dict[str, float]] = None
+    evidence: Optional[List[EvidenceItem]] = None
+    justification: Optional[str] = None

data/raw/train.json

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

requirements.txt

@@ -1,23 +1,14 @@
-pandas==2.3.0
-numpy==2.3.1
-
-scikit-learn==1.7.0
-
-xgboost==3.0.2
-lightgbm==4.6.0
-
-optuna==4.4.0
-mlflow==3.1.1
-
-matplotlib==3.10.3
-seaborn==0.13.2
-
-jupyter==1.1.1
-notebook==7.4.4
-
 fastapi==0.116.0
 uvicorn==0.35.0
 aiofiles==24.1.0
-
 joblib==1.5.1
-pytest
+scikit-learn==1.7.0
+numpy==2.3.1
+lightgbm==4.6.0
+xgboost==3.0.2
+pandas==2.3.0
+openai>=1.30.0
+python-dotenv>=1.0.0
+psycopg2-binary>=2.9.9
+pgvector>=0.3.0
+scipy>=1.13.0

start.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+set -e
+
+service postgresql start
+
+for i in $(seq 1 20); do
+    if pg_isready -U postgres -q; then
+        echo "PostgreSQL ready."
+        break
+    fi
+    echo "Waiting for PostgreSQL... ($i/20)"
+    sleep 1
+done
+
+su -c "psql -c \"DO \$\$ BEGIN IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname='mlflow') THEN CREATE USER mlflow WITH PASSWORD 'mlflow123'; END IF; END \$\$;\"" postgres
+su -c "psql -c \"SELECT 'CREATE DATABASE mlflow_db OWNER mlflow' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname='mlflow_db')\gexec\"" postgres
+su -c "psql -d mlflow_db -c \"CREATE EXTENSION IF NOT EXISTS vector;\"" postgres
+su -c "psql -d mlflow_db -c \"CREATE TABLE IF NOT EXISTS rag_claims (id SERIAL PRIMARY KEY, statement TEXT, label TEXT, reason TEXT, embedding vector(512));\"" postgres
+
+export DATABASE_URL="postgresql://mlflow:mlflow123@localhost:5432/mlflow_db"
+
+exec uvicorn api.main:app --host 0.0.0.0 --port 7860