app.py

from datasets import load_dataset
from langchain.text_splitter import RecursiveCharacterTextSplitter
from transformers import pipeline
import torch
from sentence_transformers import SentenceTransformer
import numpy as np
from tqdm import tqdm
import faiss

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

# Login using e.g. `huggingface-cli login` to access this dataset
ds = load_dataset("artemss/hotwash-rag")


# On suppose qu'il y a une colonne "text"
texts = ds["train"]["text"]
print(f"{len(texts)} documents chargés.")


# Créer le découpeur
splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,      # taille des chunks
    chunk_overlap=100,    # recouvrement pour ne pas perdre le contexte
    separators=["\n\n", "\n", ".", " "]
)

# Appliquer sur tous les textes
docs = []
for t in texts:
    chunks = splitter.split_text(t)
    docs.extend(chunks)

print(f"{len(docs)} chunks générés.")


# Ton modèle d'embedding
model_name = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
embedder = SentenceTransformer(model_name)

# Supposons que tu aies déjà tes chunks dans une liste `chunks`
# (ex : issus du chunking précédent)
# Exemple :
# chunks = ["Texte 1 ...", "Texte 2 ...", "Texte 3 ..."]

# Génération des embeddings
embeddings = embedder.encode(chunks, show_progress_bar=True, convert_to_numpy=True)

print(f"{len(embeddings)} embeddings créés. Taille de vecteur : {embeddings.shape[1]}")

# Dimension des vecteurs
dim = embeddings.shape[1]

# Création d’un index FAISS
index = faiss.IndexFlatL2(dim)
index.add(np.array(embeddings))

print(f"✅ Index FAISS créé avec {index.ntotal} vecteurs.")

def retrieve_similar_chunks(query: str, embedder, index, chunks, top_k: int = 5):
    """
    Recherche les chunks les plus pertinents pour une requête donnée.
    
    Args:
        query (str): La requête utilisateur (question, mot-clé, etc.)
        embedder: Le modèle SentenceTransformer utilisé pour les embeddings
        index: L’index FAISS contenant les vecteurs des chunks
        chunks (list): La liste des textes (chunks)
        top_k (int): Le nombre de résultats à renvoyer
        
    Returns:
        list[dict]: Une liste de résultats sous forme :
                    [{'chunk': texte, 'score': distance}, ...]
    """
    # Transformer la requête en vecteur
    query_vec = embedder.encode([query])
    
    # Rechercher les K chunks les plus proches
    distances, indices = index.search(query_vec, k=top_k)
    
    # Construire les résultats lisibles
    results = []
    for dist, idx in zip(distances[0], indices[0]):
        results.append({
            "chunk": chunks[idx],
            "score": float(dist)
        })
    
    return results


# Exemple d’utilisation :
query = "résumé sur la cybersécurité des systèmes industriels"
results = retrieve_similar_chunks(query, embedder, index, docs, top_k=3)

print("\n🔎 Résultats les plus pertinents :\n")
for r in results:
    print(f"Score: {r['score']:.4f}\nTexte: {r['chunk'][:300]}...\n")