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

print("Downloading model...")

MODEL_PATH = hf_hub_download(
    repo_id="Daffaadityp/AxonAI-MX4-2.0-GGUF",
    filename="AxonAI-MX4-2.0-Q4_K_M.gguf"
)

print("Loading model...")

llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=2048,
    n_threads=2,
    n_batch=128,
    verbose=False,
)

print("Model loaded!")

def chat(message, history):

    result = llm(
        message,
        max_tokens=128,
        temperature=0.7,
        stop=["<|im_end|>", "</s>"]
    )

    return result["choices"][0]["text"]

demo = gr.ChatInterface(
    fn=chat,
    title="AxonAI MX4 2.0",
    description="AxonAI-MX4-2.0-Q4_K_M.gguf"
)

demo.launch(server_name="0.0.0.0")