volAI_Avril / RAG_OpenAI.py
OutrageousOtter's picture
Update RAG_OpenAI.py
a8270a2 verified
raw
history blame
3.51 kB
import os
import numpy as np
import fitz # PyMuPDF pour extraction PDF
import faiss
import openai
from openai import OpenAI
from sklearn.manifold import TSNE
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from dotenv import load_dotenv
# Charger les variables d'environnement
load_dotenv()
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
# 📌 Initialisation du client OpenAI
client = OpenAI(api_key='') #todo : mettre la clé en secret et get avec os env
model_embedding = "text-embedding-ada-002"
model_chat = "gpt-4-turbo"
# 📌 Paramètres de segmentation
chunk_size = 256
chunk_overlap = 10
# 📌 Extraction et segmentation des PDF
def extract_and_chunk_pdfs(pdf_folder):
"""Extrait et segmente les textes des PDF en chunks optimisés pour OpenAI."""
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]
response = client.embeddings.create(
input=batch,
model=model_embedding
)
batch_embeddings = [data.embedding for data in 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."""
response = client.embeddings.create(
input=[question],
model=model_embedding
)
question_embedding = np.array(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 OpenAI
def generate_response(context, question):
"""Génère une réponse basée sur le contexte extrait du corpus."""
messages = [
{"role": "system", "content": f"Voici des informations contextuelles : {context}"},
{"role": "user", "content": question}
]
response = client.chat.completions.create(
model=model_chat,
messages=messages
)
return response.choices[0].message.content
# 📌 Exécuter une requête utilisateur
user_question = "Quelles souches de poulet et poules se trouvent dans ce corpus de texte ?"
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 OpenAI :")
print(answer)
# 💾 Sauvegarde des résultats
with open("openai_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 'openai_response.txt'")