from transformers import pipeline
import gradio as gr
import json

# Initialize the pipeline with the new model
pipe = pipeline("text-generation", model="Blexus/Quble_test_model_v1_INSTRUCT_v1")

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):
    # Format prompt according to the new template
    prompt = "SYSTEM: You are a helpful assistant, with no access to external functions.\n<|endofsystem|>\n"
    for user_prompt, bot_response in history:
        prompt += f"USER: {user_prompt}\n\n\nASSISTANT: {bot_response}<|endoftext|>\n"
    prompt += f"USER: {message}\n\n\nASSISTANT:"
    return prompt

def generate(
    prompt, history, temperature=0.9, max_new_tokens=4096, top_p=0.9, repetition_penalty=1.2,
):
    database = load_database()  # Load the database
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    formatted_prompt = format_prompt(prompt, history)
    if formatted_prompt in database:
        response = database[formatted_prompt]
    else:
        # Use the pipeline to generate the response
        response = pipe(formatted_prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty)[0]["generated_text"]
        response_text = response.split("ASSISTANT:")[1].strip()  # Extract the assistant's response
        database[formatted_prompt] = response_text
        save_database(database)  # Save the updated database

    yield response_text

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

with gr.Blocks(css=css) as demo:
    gr.ChatInterface(
        generate,
    )

demo.launch(debug=True)