sae-domain-adaptation

SAE Domain Adaptation v4 — Explaining EDF Document Embeddings

Version actuelle : v4 — Pipeline dual avec support Matryoshka Representation Learning via F2LLM-v2-80M.

---

Architecture v4 (Dual Pipeline)

Pipeline 1 : Analyse documentaire via LLM (Gemma-3)

Emails EDF (TSV)
  → Gemma-3 Layer 12 [d=640]
  → Max-Pool par document
  → SAE pré-entraîné Gemma Scope 16K
  → UMAP 2D + HDBSCAN clusters
  → Visualisation interactive Plotly

Avantages : Pas d'entraînement, exploitation directe des connaissances du LLM pré-entraîné.

Pipeline 2 : SAE phrase-level via F2LLM-v2-80M (Matryoshka)

F2LLM-v2-80M est un modèle Matryoshka Representation Learning : hidden_size=320, tronçable à [32, 64, 128, 256, 320].

Données entraînement (Wikipedia/FineWeb-2):
  Énergie électrique + Sports (deux domaines indépendants)
    → Découpage en phrases
    → F2LLM-v2-80M Mean-Pool [d=320] → truncation optionnelle à 256
    → SAE entraîné from-scratch (d_sae=8192, k=16)
    → Sauvegarde des poids

Emails EDF (inférence seule):
    → Découpage en phrases
    → F2LLM-v2-80M embeddings (même truncation que l'entraînement)
    → SAE chargé (poids pré-entraînés)
    → Max-Pool des activations SAE par document
    → UMAP 2D + HDBSCAN clusters

Paramètres Matryoshka :

export EMB_MODEL="codefuse-ai/F2LLM-v2-80M"      # ou F2LLM-v2-0.6B, F2LLM-v2-4B...
export MATRYOSHKA_DIM=320                           # 256 ou 320 (pour 80M)

Avantages :

• SAE entraîné sur ~20K phrases (énergie + sports) = volume suffisant
• Les emails EDF (3000 docs) sont utilisés UNIQUEMENT pour la visualisation
• Matryoshka : on peut réduire la dimension à 256 si le SAE a du mal à converger, ou monter à 320 pour plus de richesse

---

Format de données attendu

Entraînement (Wikipedia / FineWeb-2)

Automatique — pas de fichier requis. Le script stream depuis HuggingFace. Variable d'environnement pour FineWeb-2 :

export USE_FINEWEB2=true  # Streaming depuis HuggingFaceFW/fineweb-2

Inférence (EDF Emails)

Fichier TSV avec colonnes :

"Bonjour, je conteste ma facture..."	"[{\"annotation\":\"__PDL__\",\"segment\":\"1234567890123\"}]"

• document : texte complet du mail client
• segments : JSON avec annotations d'entités (PDL, personnes, etc.)

Variable d'environnement :

export EDF_EMAILS_PATH=./Mails.tsv

---

Démarrage rapide

# Dépendances
pip install torch transformers datasets sae-lens
pip install umap-learn plotly hdbscan scikit-learn pandas tqdm

# Configuration Matryoshka (optionnel, défauts=320)
export EMB_MODEL="codefuse-ai/F2LLM-v2-80M"
export MATRYOSHKA_DIM=320

# Exécution des deux pipelines
python sae_domain_adaptation_v4.py

# Résultats dans ./results_v4/ :
#   - umap_pipeline1_llm_maxpool.html
#   - umap_pipeline2_f2llm_phrases.html

---

Contenu du dépôt

---

Visualisation interactive

Les fichiers HTML générés contiennent :

• Projection UMAP 2D (métrique cosine)
• Clusters HDBSCAN colorés
• Hover tooltip : aperçu du document + top-3 features SAE activées
• Labels : annotations d'entités EDF (PDL, etc.)

---

Références

• F2LLM-v2-80M : https://huggingface.co/codefuse-ai/F2LLM-v2-80M (Matryoshka, 320-dim)
• F2LLM-v2 paper : arXiv:2603.19223
• Gemma Scope : https://huggingface.co/google/gemma-scope
• FineWeb-2 : https://huggingface.co/datasets/HuggingFaceFW/fineweb-2
• Koriagin et al., "SAE Boost", arXiv:2507.12990