RAG_Mistral.py

# -*- coding: utf-8 -*-
"""
Created on Mon Feb 24 15:51:34 2025

@author: MIPO10053340

C:/Users/MIPO10053340/OneDrive - Groupe Avril/Bureau/Salon_Agriculture_2024/Micka_API_Call/Docs_pdf/Docs_pdf/

"""

# -*- coding: utf-8 -*-
"""
Optimisation du RAG avec MistralAI - Embeddings en batch
"""
import os
import numpy as np
import fitz  # PyMuPDF pour extraction PDF
import faiss
import matplotlib.pyplot as plt
from mistralai import Mistral
from sklearn.manifold import TSNE
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from dotenv import load_dotenv

# Charger les variables d'environnement
load_dotenv()
MISTRAL_API_KEY = os.getenv('MISTRAL_API_KEY_static')

# 📌 Initialisation du client Mistral
client = Mistral(api_key=MISTRAL_API_KEY)
model_embedding = "mistral-embed"
model_chat = "mistral-large-latest"
temperature = 0.1  # Réduction de la température pour privilégier la RAG
probability = 0.9  # Ajustement de la probabilité pour plus de contrôle

# 📌 Paramètres de segmentation
chunk_size = 256  # Réduction du chunk size pour un meilleur contrôle du contexte
chunk_overlap = 15

# 📌 Extraction et segmentation des PDF
def extract_and_chunk_pdfs(pdf_folder):
    """Extrait et segmente les textes des PDF en chunks optimisés pour Mistral."""
    documents = SimpleDirectoryReader(pdf_folder).load_data()
    chunked_docs = [doc.text for doc in documents]
    return chunked_docs

# 📌 Génération des embeddings par batch
def get_embeddings_in_batches(text_chunks, batch_size=5):
    """Génère les embeddings en batch pour éviter les dépassements de tokens."""
    embeddings = []
    for i in range(0, len(text_chunks), batch_size):
        batch = text_chunks[i:i + batch_size]
        embeddings_batch_response = client.embeddings.create(
            model=model_embedding,
            inputs=batch,
        )
        batch_embeddings = [data.embedding for data in embeddings_batch_response.data]
        embeddings.extend(batch_embeddings)
    
    return np.array(embeddings).astype('float32')

# 📌 Chargement et embedding des documents
pdf_folder = 'C:/Users/MIPO10053340/OneDrive - Groupe Avril/Bureau/Salon_Agriculture_2024/Micka_API_Call/Docs_pdf/'
chunked_docs = extract_and_chunk_pdfs(pdf_folder)
embeddings = get_embeddings_in_batches(chunked_docs)

# 📌 Indexation des embeddings avec FAISS
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)

# 📌 Récupération des chunks les plus pertinents
def retrieve_relevant_chunks(question, k=5):
    """Recherche les chunks les plus pertinents en fonction de la similarité des embeddings."""
    question_embedding_response = client.embeddings.create(
        model=model_embedding,
        inputs=[question],
    )
    question_embedding = np.array(question_embedding_response.data[0].embedding).astype('float32').reshape(1, -1)
    distances, indices = index.search(question_embedding, k)
    return [chunked_docs[i] for i in indices[0]]

# 📌 Génération de réponse avec MistralAI
def generate_response(context, question):
    """Génère une réponse basée sur le contexte extrait du corpus avec une basse température et un contrôle de probabilité."""
    messages = [
        {"role": "system", "content": f"Voici des informations contextuelles à utiliser avec priorité : {context}"},
        {"role": "user", "content": question}
    ]
    
    response = client.chat.complete(model=model_chat, messages=messages, temperature=temperature, probability=probability)
    return response.choices[0].message.content

# 📌 Exécuter une requête utilisateur
user_question = "Quelles sont les souches de poulets ou poules présentent dans les publications de notre corpus utilisé pour la RAG"
relevant_chunks = retrieve_relevant_chunks(user_question)
context = "\n".join(relevant_chunks)
answer = generate_response(context, user_question)

# 📊 Affichage de la réponse
print("\n🔹 Réponse Mistral :")
print(answer)

# # 📊 Visualisation des embeddings avec t-SNE
# tsne = TSNE(n_components=2, perplexity=min(30, max(2, embeddings.shape[0] - 1)), random_state=42)
# embeddings_2d = tsne.fit_transform(embeddings)

# plt.figure(figsize=(10, 8))
# plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], alpha=0.5)
# plt.title('Visualisation des embeddings avec t-SNE')
# plt.xlabel('Dimension 1')
# plt.ylabel('Dimension 2')
# plt.show()

# 💾 Sauvegarde des résultats
with open("mistral_response.txt", "w", encoding="utf-8") as f:
    f.write(f"Question : {user_question}\n")
    f.write(f"Réponse :\n{answer}\n")

print("\n✅ Réponse enregistrée dans 'mistral_response.txt'")