import numpy as np
import faiss
import pandas as pd
from transformers import pipeline
import gradio as gr

# Charger les embeddings et l'index FAISS
def load_embeddings_and_index(embeddings_file, index_file):
    embeddings = np.load(embeddings_file)
    index = faiss.read_index(index_file)
    return embeddings, index

# Charger les données sources
def load_data(file_path):
    data = pd.read_excel(file_path)
    data['combined'] = data[['Product_name', 'price']].apply(
        lambda row: f"Product: {row['Product_name']} | Price: {row['price']}", axis=1
    )
    return data

# Effectuer une recherche dans l'index FAISS
def query_faiss_index(query, embeddings, index, data, embedding_model, top_k=5):
    query_embedding = embedding_model.encode([query])
    distances, indices = index.search(query_embedding, top_k)
    results = [data['combined'].iloc[idx] for idx in indices[0]]
    return results

# Générer une réponse avec le modèle LLM
def generate_answer(query, context, llm_pipeline):
    prompt = f"Given the following context, answer the question:\n\nContext:\n{context}\n\nQuestion: {query}\nAnswer:"
    response = llm_pipeline(prompt, max_length=500, num_return_sequences=1)
    return response[0]['generated_text']

# Interface principale
def main(embeddings_file, index_file, data_file):
    # Charger les données et les fichiers nécessaires
    embeddings, index = load_embeddings_and_index(embeddings_file, index_file)
    data = load_data(data_file)

    # Charger le modèle LLM
    llm_pipeline = pipeline("text-generation", model="HuggingFaceH4/zypher-llm")

    # Interface utilisateur avec Gradio
    def gradio_interface(query):
        # Recherche dans FAISS
        context = query_faiss_index(query, embeddings, index, data, embedding_model, top_k=5)
        context_text = "\n".join(context)
        # Générer une réponse à l'aide du LLM
        answer = generate_answer(query, context_text, llm_pipeline)
        return answer

    # Lancer l'application Gradio
    interface = gr.Interface(fn=gradio_interface, inputs="text", outputs="text", title="RAG Chatbot")
    interface.launch()

# Exemple d'exécution
if __name__ == "__main__":
    # Fichiers nécessaires
    embeddings_file = "embeddings.npy"
    index_file = "faiss_index.bin"
    data_file = "avito.xls"

    # Charger le modèle d'embedding
    from sentence_transformers import SentenceTransformer
    embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

    # Lancer le programme
    main(embeddings_file, index_file, data_file)