tlavandier/OpenData-Bordeaux-RSE

1import streamlit as st2import openai3import os4from dotenv import load_dotenv5from PyPDF2 import PdfReader6import faiss7import numpy as np8from sentence_transformers import SentenceTransformer9 10# Charger les variables d'environnement11load_dotenv(".env")12YOUR_API_KEY = os.getenv("API_TOKEN_PERPLEXITYAI")13openai.api_key = YOUR_API_KEY14 15# Charger le modèle pour les embeddings16embedding_model = SentenceTransformer('all-MiniLM-L6-v2')17 18def extract_text_from_pdf(file):19    """Extraire le texte d'un fichier PDF."""20    reader = PdfReader(file)21    text = ""22    for page in reader.pages:23        text += page.extract_text()24    return text25 26def create_embeddings(text):27    """Créer des embeddings pour un texte donné."""28    return embedding_model.encode([text])[0]29 30def create_vector_index(texts):31    """Créer un index vectoriel à partir des textes donnés."""32    embeddings = [create_embeddings(text) for text in texts]33    dimension = len(embeddings[0])34    index = faiss.IndexFlatL2(dimension)35    index.add(np.array(embeddings).astype(np.float32))36    return index, texts37 38def create_pdf_vector_store(pdf_files):39    """Créer le `vector store` à partir des fichiers PDF téléchargés."""40    texts = [extract_text_from_pdf(file) for file in pdf_files]41    return create_vector_index(texts)42 43def search_vector_store(query, index, texts, k=3):44    """Effectuer une recherche dans le `vector store`."""45    query_embedding = create_embeddings(query)46    _, indices = index.search(np.array([query_embedding]).astype(np.float32), k)47    return [texts[i] for i in indices[0]]48 49def perform_chat_with_context(query, index, texts, model):50    """Générer une réponse basée sur les documents les plus pertinents."""51    context = "\n".join(search_vector_store(query, index, texts))52    messages = [53        {"role": "system", "content": "Vous êtes un assistant IA."},54        {"role": "user", "content": f"Contexte : {context}\n\nQuestion : {query}"}55    ]56    client = openai.OpenAI(api_key=YOUR_API_KEY, base_url="https://api.perplexity.ai")57    response = client.chat.completions.create(model=model, messages=messages)58    st.write(response.choices[0].message.content)59 60def display_transition_plan():61    """Interface Streamlit pour définir un plan de transition."""62    st.title("Définir un plan de transition pour une marque")63    brand = st.text_input("Nom de la marque")64    uploaded_files = st.file_uploader("Téléchargez les documents PDF", type="pdf", accept_multiple_files=True)65    66    model = st.selectbox("Choisissez un modèle d'IA :", ["sonar-medium-online", "mistral-7b-instruct", "mixtral-8x7b-instruct", "sonar-small-online", "codellama-70b-instruct"])67    68    if st.button("Proposer un plan de transition"):69        if brand and uploaded_files:70            index, texts = create_pdf_vector_store(uploaded_files)71            query = f"Proposez un plan de transition pour la marque {brand}"72            perform_chat_with_context(query, index, texts, model)73        else:74            st.write("Veuillez saisir le nom de la marque et télécharger les documents PDF.")75