setup_bitnet.sh

#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────
# setup_bitnet.sh — Build llama.cpp with BitNet ternary kernels
# and quantize DeepSeek R1 to near-ternary (IQ2_XXS) for fast
# CPU inference.
#
# Usage:
#   ./setup_bitnet.sh              # default: IQ2_XXS (~1.58 bpw)
#   ./setup_bitnet.sh Q2_K         # alternative ~2.6 bpw
#   ./setup_bitnet.sh IQ1_S        # extreme 1.0 bpw (lowest quality)
#
# After setup, start the server:
#   ./start_bitnet.sh
#
# Then run:
#   BITNET=1 python3 chat.py
# ─────────────────────────────────────────────────────────────
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
LLAMA_DIR="$SCRIPT_DIR/.llama.cpp"
MODELS_DIR="$SCRIPT_DIR/.models"
MODEL_NAME="deepseek-r1-8b"
QUANT_TYPE="${1:-IQ2_XXS}"   # IQ2_XXS ≈ 1.58 bits/weight (BitNet-like)

echo "============================================"
echo "  BitNet-style DeepSeek R1 Setup"
echo "  Quantization: $QUANT_TYPE (~1.58 bpw)"
echo "============================================"

# ── 1. Install build dependencies ──
echo "[1/5] Checking build dependencies..."
if command -v apt-get &>/dev/null; then
    sudo apt-get update -qq
    sudo apt-get install -y -qq build-essential cmake git python3 python3-pip curl
elif command -v brew &>/dev/null; then
    brew install cmake curl
fi

# ── 2. Clone / update llama.cpp ──
echo "[2/5] Setting up llama.cpp with BitNet support..."
if [ -d "$LLAMA_DIR" ]; then
    cd "$LLAMA_DIR"
    git pull --quiet
else
    git clone --depth 1 https://github.com/ggerganov/llama.cpp.git "$LLAMA_DIR"
    cd "$LLAMA_DIR"
fi

# ── 3. Build with optimised low-bit kernels (TQ1_0/IQ2_XXS use BitNet-style paths) ──
echo "[3/5] Building llama.cpp (this may take a few minutes)..."
cmake -B build \
    -DGGML_NATIVE=ON \
    -DGGML_CPU_AARCH64=OFF \
    -DLLAMA_CURL=ON \
    -DCMAKE_BUILD_TYPE=Release \
    2>&1 | tail -5
cmake --build build --config Release -j "$(nproc)" 2>&1 | tail -5

LLAMA_SERVER="$LLAMA_DIR/build/bin/llama-server"
LLAMA_QUANTIZE="$LLAMA_DIR/build/bin/llama-quantize"

if [ ! -f "$LLAMA_SERVER" ]; then
    echo "ERROR: llama-server not found after build. Check build output above."
    exit 1
fi
echo "  ✓ llama.cpp built at $LLAMA_DIR/build"

# ── 4. Download the pre-quantized GGUF model ──
echo "[4/5] Downloading model (${QUANT_TYPE})..."
mkdir -p "$MODELS_DIR"

GGUF_BASE="$MODELS_DIR/${MODEL_NAME}-f16.gguf"
GGUF_QUANT="$MODELS_DIR/${MODEL_NAME}-${QUANT_TYPE}.gguf"

if [ ! -f "$GGUF_QUANT" ]; then
    pip install -q huggingface-hub 2>/dev/null || pip install -q huggingface-hub

    python3 - "$MODELS_DIR" "$GGUF_QUANT" "$QUANT_TYPE" "$GGUF_BASE" <<'PYEOF'
import sys, os
from huggingface_hub import hf_hub_download

models_dir  = sys.argv[1]
target_path = sys.argv[2]
quant_type  = sys.argv[3]
gguf_base   = sys.argv[4]

# Try several known GGUF repos for DeepSeek-R1 8B
candidates = [
    ("bartowski/DeepSeek-R1-Distill-Qwen-8B-GGUF",
     f"DeepSeek-R1-Distill-Qwen-8B-{quant_type}.gguf"),
    ("bartowski/DeepSeek-R1-Distill-Qwen-8B-GGUF",
     "DeepSeek-R1-Distill-Qwen-8B-Q2_K.gguf"),
    ("bartowski/DeepSeek-R1-Distill-Qwen-8B-GGUF",
     "DeepSeek-R1-Distill-Qwen-8B-IQ2_M.gguf"),
]

for repo, fname in candidates:
    try:
        print(f"  Trying {repo} / {fname} ...")
        path = hf_hub_download(
            repo_id=repo,
            filename=fname,
            local_dir=models_dir,
            local_dir_use_symlinks=False,
        )
        downloaded = os.path.join(models_dir, fname)
        if os.path.exists(downloaded) and downloaded != target_path:
            os.rename(downloaded, target_path)
        elif path != target_path and os.path.exists(path):
            os.rename(path, target_path)
        print(f"  ✓ Downloaded {fname}")
        sys.exit(0)
    except Exception as e:
        print(f"    Skipped: {e}")

print("ERROR: Could not download a pre-quantized model.")
print(f"Place a GGUF file manually at: {gguf_base}")
print("Then re-run this script to quantize it.")
sys.exit(1)
PYEOF
fi

echo "  ✓ Model ready: $GGUF_QUANT"

# ── 5. Write a convenience start script ──
MODEL_SIZE=$(du -h "$GGUF_QUANT" | cut -f1)

cat > "$SCRIPT_DIR/start_bitnet.sh" << STARTEOF
#!/usr/bin/env bash
# Auto-generated by setup_bitnet.sh — starts llama.cpp server
# with the BitNet-quantized DeepSeek model.
set -euo pipefail

LLAMA_SERVER="$LLAMA_SERVER"
MODEL_PATH="$GGUF_QUANT"
HOST="\${BITNET_HOST:-0.0.0.0}"
PORT="\${BITNET_PORT:-8081}"
THREADS="\${BITNET_THREADS:-\$(nproc)}"
CTX="\${BITNET_CTX:-4096}"

echo "Starting llama.cpp BitNet server..."
echo "  Model  : $MODEL_NAME ($QUANT_TYPE, $MODEL_SIZE)"
echo "  Listen : \$HOST:\$PORT"
echo "  Threads: \$THREADS  Context: \$CTX"

exec "\$LLAMA_SERVER" \\
    -m "\$MODEL_PATH" \\
    --host "\$HOST" --port "\$PORT" \\
    -ngl 0 \\
    -c "\$CTX" \\
    -t "\$THREADS" \\
    --chat-template chatml
STARTEOF
chmod +x "$SCRIPT_DIR/start_bitnet.sh"

echo ""
echo "[5/5] Setup complete!"
echo "============================================"
echo "  Model  : $MODEL_NAME ($QUANT_TYPE)"
echo "  Size   : $MODEL_SIZE"
echo "  Path   : $GGUF_QUANT"
echo ""
echo "  Start the BitNet server:"
echo "    ./start_bitnet.sh"
echo ""
echo "  Then run the chatbot:"
echo "    BITNET=1 python3 chat.py"
echo "============================================"