darkbert_model/darkbert_final.py

import os
from sentence_transformers import SentenceTransformer
from pathlib import Path
import numpy as np
import logging

# Configuration
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
logging.basicConfig(level=logging.INFO)

def main():
    try:
        # 1. Charger le modèle
        logging.info("Chargement du modèle...")
        model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
        
        # 2. Lire les fichiers
        logging.info("Lecture des documents...")
        docs = []
        for file in Path('input').glob('*.html'):
            with open(file, 'r', encoding='utf-8') as f:
                docs.append(f.read())
        
        # 3. Générer les embeddings
        logging.info("Création des embeddings...")
        embeddings = model.encode(docs, show_progress_bar=True)
        
        # 4. Sauvegarder
        np.save('output/embeddings.npy', embeddings)
        logging.info(f"Terminé ! Résultats sauvegardés dans output/")
        
    except Exception as e:
        logging.error(f"Erreur : {str(e)}")

if __name__ == "__main__":
    main()