import gradio as gr
from huggingface_hub import InferenceClient
from transformers import RagTokenizer, RagTokenForGeneration
from typing import List, Dict, Tuple
import re
import os
import torch
from math import ceil

# Load the RAG model and tokenizer
rag_tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
rag_model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq")

# Load your PDF document
pdf_path = "apexcustoms.pdf"
with open(pdf_path, 'rb') as f:
    pdf_text = f.read().decode('utf-8', errors='ignore')

# Split the PDF text into chunks
split_pattern = r'\n\n'
doc_chunks = re.split(split_pattern, pdf_text)

# Preprocess the corpus
corpus = rag_tokenizer(doc_chunks, return_tensors="pt", padding=True, truncation=True).input_ids

# Pad the corpus to be a multiple of `n_docs`
n_docs = rag_model.config.n_docs
corpus_length = corpus.size(-1)
pad_length = ceil(corpus_length / n_docs) * n_docs - corpus_length
corpus = torch.nn.functional.pad(corpus, (0, pad_length), mode='constant', value=rag_tokenizer.pad_token_id)

"""
For more information on huggingface_hub Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")


def respond(
    message,
    history: List[Tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_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})

    # Tokenize the input
    inputs = rag_tokenizer(message, return_tensors="pt")
    input_ids = inputs.pop("input_ids")

    # Generate with the RAG model
    output_ids = rag_model.generate(
        input_ids=input_ids,
        context_input_ids=corpus,
        max_length=max_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=0,
        num_beams=2,
    )
    retrieved_context = rag_tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
        context=retrieved_context,  # Include the retrieved context
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
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.7, 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()