app.py

import os
import subprocess
import streamlit as st
from huggingface_hub import snapshot_download

def check_directory_path(directory_name: str) -> str:
    if os.path.exists(directory_name):
        path = os.path.abspath(directory_name)
        return str(path)

# Define quantization types
QUANT_TYPES = [
    "Q2_K", "Q3_K_M", "Q3_K_S", "Q4_K_M", "Q4_K_S",
    "Q5_K_M", "Q5_K_S", "Q6_K"
]

model_dir_path = check_directory_path("/app/llama.cpp")

def download_model(hf_model_name, output_dir="/tmp/models"):
    """
    Downloads a Hugging Face model and saves it locally.
    """
    st.write(f"📥 Downloading `{hf_model_name}` from Hugging Face...")
    os.makedirs(output_dir, exist_ok=True)
    snapshot_download(repo_id=hf_model_name, local_dir=output_dir, local_dir_use_symlinks=False)
    st.success("✅ Model downloaded successfully!")

def convert_to_gguf(model_dir, output_file):
    """
    Converts a Hugging Face model to GGUF format.
    """
    st.write(f"🔄 Converting `{model_dir}` to GGUF format...")
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    st.write(model_dir_path)
    cmd = [
        "python3", "/app/llama.cpp/convert_hf_to_gguf.py", model_dir,
        "--outtype", "f16", "--outfile", output_file
    ]
    process = subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if process.returncode == 0:
        st.success(f"✅ Conversion complete: `{output_file}`")
    else:
        st.error(f"❌ Conversion failed: {process.stderr}")

def quantize_llama(model_path, quantized_output_path, quant_type):
    """
    Quantizes a GGUF model.
    """
    st.write(f"⚡ Quantizing `{model_path}` with `{quant_type}` precision...")
    os.makedirs(os.path.dirname(quantized_output_path), exist_ok=True)
    quantize_path = "/app/llama.cpp/build/bin/llama-quantize"
    
    cmd = [
        "/app/llama.cpp/build/bin/llama-quantize", 
        model_path, 
        quantized_output_path,
        quant_type
    ]
    
    process = subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    if process.returncode == 0:
        st.success(f"✅ Quantized model saved at `{quantized_output_path}`")
    else:
        st.error(f"❌ Quantization failed: {process.stderr}")

def automate_llama_quantization(hf_model_name, quant_type):
    """
    Orchestrates the entire quantization process.
    """
    output_dir = "/tmp/models"
    gguf_file = os.path.join(output_dir, f"{hf_model_name.replace('/', '_')}.gguf")
    quantized_file = gguf_file.replace(".gguf", f"-{quant_type}.gguf")

    progress_bar = st.progress(0)

    # Step 1: Download
    st.write("### Step 1: Downloading Model")
    download_model(hf_model_name, output_dir)
    progress_bar.progress(33)

    # Step 2: Convert to GGUF
    st.write("### Step 2: Converting Model to GGUF Format")
    convert_to_gguf(output_dir, gguf_file)
    progress_bar.progress(66)

    # Step 3: Quantize Model
    st.write("### Step 3: Quantizing Model")
    quantize_llama(gguf_file, quantized_file, quant_type.lower())
    progress_bar.progress(100)

    st.success(f"🎉 All steps completed! Quantized model available at: `{quantized_file}`")
    return quantized_file

# Streamlit UI
st.title("🦙 LLaMA Model Quantization (llama.cpp)")

hf_model_name = st.text_input("Enter Hugging Face Model Name", "Qwen/Qwen2.5-1.5B")
quant_type = st.selectbox("Select Quantization Type", QUANT_TYPES)
start_button = st.button("🚀 Start Quantization")

if start_button:
    with st.spinner("Processing..."):
        quantized_model_path = automate_llama_quantization(hf_model_name, quant_type)
    if quantized_model_path:
        with open(quantized_model_path, "rb") as f:
            st.download_button("⬇️ Download Quantized Model", f, file_name=os.path.basename(quantized_model_path))