from huggingface_hub import InferenceClient
import gradio as gr
import json

client = InferenceClient(
    "mistralai/Mistral-7B-Instruct-v0.1"
)

DATABASE_PATH = "database.json"

def load_database():
    try:
        with open(DATABASE_PATH, "r") as file:
            return json.load(file)
    except FileNotFoundError:
        return {}

def save_database(database):
    with open(DATABASE_PATH, "w") as file:
        json.dump(database, file)

def format_prompt(message, history):
    prompt = "<s>"
    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST]"
        prompt += f" {bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"
    return prompt

def generate_response(prompt, database):
    if prompt in database:
        return database[prompt]
    else:
        response = next(client.text_generation(prompt, details=True, return_full_text=False)).token.text
        database[prompt] = response
        save_database(database)
        return response

def generate(
    prompt, history, database, temperature=0.9, max_new_tokens=2000, top_p=0.9, repetition_penalty=1.2,
):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    formatted_prompt = format_prompt(prompt, history)
    response = generate_response(formatted_prompt, database)
    yield response

database = load_database()

css = """
  #mkd {
    height: 500px; 
    overflow: auto; 
    border: 1px solid #ccc; 
  }
"""

with gr.Blocks(css=css) as demo:
    gr.ChatInterface(
        generate,
        examples=[["What is the secret to life?"], ["Write me a recipe for pancakes."], ["Write a short story about Paris."]],
        database=database
    )

demo.launch(debug=True)