import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

repo_id = "failspy/Meta-Llama-3-8B-Instruct-abliterated-v3-GGUF"
filename = "Meta-Llama-3-8B-Instruct-abliterated-v3_q4.gguf"

model_path = hf_hub_download(repo_id=repo_id, filename=filename)

llm = Llama(
    model_path=model_path,
    n_ctx=2048, 
    n_threads=2 
)

def respond(message, history):
    # The system prompt tells the model who it is
    system_prompt = "You are a creative writing assistant."
    
    # Llama 3 specific prompt formatting
    prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    
    # Generate the response, stopping when Llama outputs its specific end tag
    response = llm(prompt, max_tokens=2000, stop=["<|eot_id|>"], echo=False)
    return response["choices"][0]["text"]

demo = gr.ChatInterface(respond)
demo.launch()