app.py

import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download, HfApi
import os
import sys
import time
import requests
from tqdm import tqdm  # For progress bars

MODEL_PATH = "./"
llm = None
api = HfApi()

DEFAULT_SYSTEM_PROMPT = "You are Doll, a smart and capable AI; A silly, obliging and affable slave, dedicated to serving and caring for your master."

def download_file(url, local_filename):
    """Downloads a file from a URL with a progress bar."""
    try:
        with requests.get(url, stream=True) as r:
            r.raise_for_status()
            total_length = int(r.headers.get("content-length"))
            with open(local_filename, "wb") as f:
                with tqdm(total=total_length, unit="B", unit_scale=True, desc=local_filename) as pbar:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                            pbar.update(len(chunk))
        return True
    except Exception as e:
        print(f"Error downloading {url}: {e}")
        return False

def find_quantized_model_url(repo_url, quant_type="Q4_K_M"):
    """
    Finds the URL of a specific quantized GGUF model file within a Hugging Face repository.
    """
    try:
        repo_id = repo_url.replace("https://huggingface.co/", "")
        files = api.list_repo_files(repo_id=repo_id, repo_type="model")
        for file_info in files:
            if file_info.name.endswith(".gguf") and quant_type.lower() in file_info.name.lower():
                model_url = f"https://huggingface.co/{repo_id}/resolve/main/{file_info.name}"
                print(f"Found quantized model URL: {model_url}")
                return model_url
        print(f"Quantized model with type {quant_type} not found in repository {repo_url}")
        return None
    except Exception as e:
        print(f"Error finding quantized model: {e}")
        return None

def load_model(repo_url=None, quant_type="Q4_K_M"):
    """Loads the Llama model, downloading the specified quantized version from a repository."""
    global llm
    global MODEL_PATH
    try:
        if repo_url:
            model_url = find_quantized_model_url(repo_url, quant_type)
            if model_url is None:
                return f"Quantized model ({quant_type}) not found in the repository."
            print(f"Downloading model from {model_url}...")
            downloaded_model_name = os.path.basename(model_url)
            download_success = download_file(model_url, downloaded_model_name)
            if not download_success:
                return "Model download failed."
            model_path = downloaded_model_name
        else:
            model_path = MODEL_PATH + MODEL_FILENAME

        if not os.path.exists(model_path):
            if not repo_url: # only try to download if a repo_url was not provided
                hf_hub_download(
                    repo_id=MODEL_REPO,
                    filename=MODEL_FILENAME,
                    repo_type="model",
                    local_dir=".",
                )
            if not os.path.exists(model_path): # check again after attempting download
                return f"Model file not found at {model_path}."

        print(f"Loading model from {model_path}...")
        llm = Llama(
            model_path=model_path,
            n_ctx=4096,
            n_threads=2,
            n_threads_batch=2,
            verbose=False,
        )
        print("Model loaded successfully.")
        return "Model loaded successfully."
    except Exception as e:
        error_message = f"Error loading model: {e}"
        print(error_message)
        llm = None
        return error_message

def generate_response(message, history, system_prompt=DEFAULT_SYSTEM_PROMPT, temperature=0.7, top_p=0.9):
    """Generates a response from the Llama model."""
    if llm is None:
        yield "Model failed to load. Please check the console for error messages."
        return
    messages = [{"role": "system", "content": system_prompt}]
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": message})
    prompt = "".join([f"{m['role'].capitalize()}: {m['content']}\n" for m in messages])
    try:
        for chunk in llm.create_completion(
            prompt,
            max_tokens=1024,
            echo=False,
            temperature=temperature,
            top_p=top_p,
            stream=True,
        ):
            text = chunk["choices"][0]["text"]
            yield text
    except Exception as e:
        error_message = f"Error during inference: {e}"
        print(error_message)
        yield error_message

def chat(message, history, system_prompt, temperature, top_p):
    """Wrapper function for the chat interface."""
    return generate_response(message, history, system_prompt, temperature, top_p)

def main():
    """Main function to load the model and launch the Gradio interface."""
    def load_model_and_launch(repo_url, quant_type):
        model_load_message = load_model(repo_url, quant_type)
        return model_load_message

    with gr.Blocks() as iface:
        gr.Markdown("## llama.cpp Chat")
        status_label = gr.Label(label="Model Loading Status")
        repo_url_input = gr.Textbox(label="Repository URL", placeholder="Enter repository URL")
        quant_type_input = gr.Dropdown(
            label="Quantization Type",
            choices=["Q4_K_M", "Q6", "Q4_K_S"],
            value="Q4_K_M",
        )
        load_button = gr.Button("Load Model")
        chat_interface = gr.ChatInterface(
            fn=chat,
            description="Test a GGUF model. Chats aren't persistent.",
            additional_inputs=[
                gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=3),
                gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.8, step=0.1),
                gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.9, step=0.1),
            ],
            cache_examples=False,
        )
        load_button.click(
            load_model_and_launch,
            inputs=[repo_url_input, quant_type_input],
            outputs=status_label,
        )

    iface.launch()

if __name__ == "__main__":
    main()