# chatbot.py

import gradio as gr
from huggingface_hub import InferenceClient
import faiss
import json
from sentence_transformers import SentenceTransformer

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Load the FAISS index and the sentence transformer model
index = faiss.read_index("apexcustoms_index.faiss")
model = SentenceTransformer('all-mpnet-base-v2')

# Load the extracted text
with open("apexcustoms.json", "r") as f:
    documents = json.load(f)

def retrieve_documents(query, k=5):
    query_embedding = model.encode([query])
    distances, indices = index.search(query_embedding, k)
    return [documents[i] for i in indices[0]]

def respond(message, history, system_message, max_tokens, temperature, top_p):
    # Retrieve relevant documents
    relevant_docs = retrieve_documents(message)
    context = "\n\n".join(relevant_docs)

    messages = [{"role": "system", "content": system_message},
                {"role": "user", "content": f"Context: {context}\n\n{message}"}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a helpful car configuration assistant, specifically you are the assistant for Apex Customs (https://www.apexcustoms.com/). Given the user's input, provide suggestions for car models, colors, and customization options. Be creative and conversational in your responses. You should remember the user car model and tailor your answers accordingly. \n\nUser: ", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.3, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
)

if __name__ == "__main__":
    demo.launch()