import gradio as gr
from gpt4all import GPT4All
from huggingface_hub import hf_hub_download

title = "Synthia 7B GGUF"
description = "This is a demo of Synthia 7B GGUF running on spaces cpu basic hardware (Free)"

model_path = "models"
model_name = "synthia-7b-v2.0-16k.Q4_K_M.gguf"
model_path = hf_hub_download(repo_id="TheBloke/SynthIA-7B-v2.0-16k-GGUF", filename=model_name, local_dir=model_path, local_dir_use_symlinks=False)

try:
    model = GPT4All(model_path, allow_download=False, device="cpu")
    print("Model loaded successfully.")
except Exception as e:
    print("Error loading the model:", str(e))

model.config["promptTemplate"] = "[INST] {0} [/INST]"
model.config["systemPrompt"] = ""
model._is_chat_session_activated = False

max_new_tokens = 512

def generater(message, history, temperature, top_p, top_k):
    prompt = "<s>"
    for user_message, assistant_message in history:
        prompt += model.config["promptTemplate"].format(user_message)
        prompt += assistant_message + "</s>"
    prompt += model.config["promptTemplate"].format(message)
    outputs = []    
    for token in model.generate(prompt=prompt, temp=temperature, top_k = top_k, top_p = top_p, max_tokens = max_new_tokens, streaming=True):
        outputs.append(token)
    return "".join(outputs)

def vote(data: gr.LikeData):
    if data.liked:
        return
    else:
        return

chatbot = gr.Chatbot(avatar_images=('resourse/user-icon.png', 'resourse/chatbot-icon.png'),bubble_full_width = False)

additional_inputs=[
    gr.Slider(
        label="temperature",
        value=0.5,
        minimum=0.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic.",
    ),
    gr.Slider(
        label="top_p",
        value=1.0,
        minimum=0.0,
        maximum=1.0,
        step=0.01,
        interactive=True,
        info="0.1 means only the tokens comprising the top 10% probability mass are considered. Suggest set to 1 and use temperature. 1 means 100% and will disable it",
    ),
    gr.Slider(
        label="top_k",
        value=40,
        minimum=0,
        maximum=1000,
        step=1,
        interactive=True,
        info="limits candidate tokens to a fixed number after sorting by probability. Setting it higher than the vocabulary size deactivates this limit.",
    )
]

character = "Sherlock Holmes"
series = "Arthur Conan Doyle's novel"

iface = gr.ChatInterface(
    fn = generater,
    title=title,
    description = description,
    chatbot=chatbot,
    additional_inputs=additional_inputs,
    examples=[
        ["Hello there! How are you doing?"],
        ["How many hours does it take a man to eat a Helicopter?"],
        ["You are a helpful and honest assistant. Always answer as helpfully as possible. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information."],
        ["I want you to act as a spoken English teacher and improver. I will speak to you in English and you will reply to me in English to practice my spoken English. I want you to strictly correct my grammar mistakes, typos, and factual errors. I want you to ask me a question in your reply. Now let's start practicing, you could ask me a question first. Remember, I want you to strictly correct my grammar mistakes, typos, and factual errors."],
        [f"I want you to act like {character} from {series}. I want you to respond and answer like {character} using the tone, manner and vocabulary {character} would use. Do not write any explanations. Only answer like {character}. You must know all of the knowledge of {character}."]
    ]
)

with gr.Blocks(css="resourse/style/custom.css") as demo:
    chatbot.like(vote, None, None)
    iface.launch()

if __name__ == "__main__":
    demo.queue(max_size=3).launch()