Upload folder using huggingface_hub

Dockerfile

@@ -1,43 +1,25 @@
 FROM python:3.11-slim
 
 ENV DEBIAN_FRONTEND=noninteractive \
-    MODEL_REPO=NANI-Nithin/K2-Horizon-0.9B-GGUF \
-    MODEL_FILE=K2-Horizon-0.9B-Q4_K_M.gguf \
-    MODEL_DIR=/data/models/k2-horizon \
-    LLAMA_VERSION=b9592 \
-    LLAMA_DIR=/opt/llama.cpp \
-    LLAMA_SERVER_BIN=/opt/llama.cpp/llama-server \
-    LD_LIBRARY_PATH=/opt/llama.cpp \
-    LLAMA_HOST=0.0.0.0 \
-    LLAMA_PORT=7860 \
-    THREADS=8 \
-    CTX_SIZE=4096 \
-    BATCH_SIZE=default \
-    UBATCH_SIZE=default \
-    CACHE_TYPE_K=default \
-    CACHE_TYPE_V=default \
-    GPU_LAYERS=0 \
-    TEMPERATURE=0.6 \
-    TOP_P=0.95 \
-    TOP_K=40 \
-    REPEAT_PENALTY=1.1 \
-    HF_XET_HIGH_PERFORMANCE=1 \
-    PYTHONUNBUFFERED=1
+    PYTHONUNBUFFERED=1 \
+    HF_HUB_ENABLE_HF_TRANSFER=0
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     ca-certificates \
     curl \
+    git \
     libgomp1 \
     libstdc++6 \
     && rm -rf /var/lib/apt/lists/*
 
-RUN mkdir -p "${LLAMA_DIR}" \
-    && curl -fL "https://github.com/ggml-org/llama.cpp/releases/download/${LLAMA_VERSION}/llama-${LLAMA_VERSION}-bin-ubuntu-x64.tar.gz" \
-        | tar -xz --strip-components=1 -C "${LLAMA_DIR}" \
-    && chmod +x "${LLAMA_SERVER_BIN}"
-
 RUN pip install --no-cache-dir \
-    huggingface_hub
+    "torch>=2.0.0" \
+    "transformers>=4.40.0" \
+    "gradio>=4.0.0" \
+    "huggingface_hub" \
+    "accelerate" \
+    "einops" \
+    "jinja2"
 
 WORKDIR /app
 COPY app.py /app/app.py

README.md

@@ -10,18 +10,18 @@ pinned: false
 models:
   - IFM/K2-Horizon-0.9B
 tags:
-  - llama.cpp
-  - gguf
   - k2-horizon
+  - transformers
+  - pytorch
+  - gradio
   - cpu
 ---
 
-# K2-Horizon-0.9B Chat
+# K2-Horizon-0.9B Chat Demo
 
-Hugging Face Space running K2-Horizon-0.9B with llama.cpp on CPU.
+Hugging Face Space running `IFM/K2-Horizon-0.9B` with Transformers + Gradio UI on CPU.
 
-- Model: `IFM/K2-Horizon-0.9B` (GGUF quant: `NANI-Nithin/K2-Horizon-0.9B-GGUF`)
-- Default quant: `K2-Horizon-0.9B-Q4_K_M.gguf`
-- Backend: `llama.cpp` `llama-server`
-- UI: native `llama.cpp` web UI
+- Model: `IFM/K2-Horizon-0.9B`
+- Backend: PyTorch / Hugging Face Transformers (`AutoModelForCausalLM`)
+- UI: Gradio ChatInterface
 - Target: HF Spaces Docker CPU

app.py

@@ -1,159 +1,128 @@
 from __future__ import annotations
 
-import json
 import os
 import sys
-import urllib.parse
-import urllib.request
-from pathlib import Path
-
-from huggingface_hub import hf_hub_download
-
-
-MODEL_REPO = os.getenv("MODEL_REPO", "NANI-Nithin/K2-Horizon-0.9B-GGUF")
-MODEL_FILE = os.getenv("MODEL_FILE", "K2-Horizon-0.9B-Q4_K_M.gguf")
-MODEL_DIR = Path(os.getenv("MODEL_DIR", "/data/models/k2-horizon"))
-CHAT_TEMPLATE_FILE = Path(os.getenv("CHAT_TEMPLATE_FILE", "/data/models/k2-horizon/chat_template.jinja"))
-
-LLAMA_SERVER_BIN = os.getenv("LLAMA_SERVER_BIN", "/opt/llama.cpp/llama-server")
-LLAMA_HOST = os.getenv("LLAMA_HOST", "0.0.0.0")
-LLAMA_PORT = os.getenv("LLAMA_PORT", "7860")
-
-THREADS = os.getenv("THREADS", "8")
-CTX_SIZE = os.getenv("CTX_SIZE", "4096")
-BATCH_SIZE = os.getenv("BATCH_SIZE", "default")
-UBATCH_SIZE = os.getenv("UBATCH_SIZE", "default")
-GPU_LAYERS = os.getenv("GPU_LAYERS", "0")
-FLASH_ATTN = os.getenv("FLASH_ATTN", "default")
-CACHE_TYPE_K = os.getenv("CACHE_TYPE_K", "default")
-CACHE_TYPE_V = os.getenv("CACHE_TYPE_V", "default")
-
-TEMPERATURE = os.getenv("TEMPERATURE", "0.6")
-TOP_P = os.getenv("TOP_P", "0.95")
-TOP_K = os.getenv("TOP_K", "40")
-REPEAT_PENALTY = os.getenv("REPEAT_PENALTY", "1.1")
-
-
-def log(message: str) -> None:
-    print(f"[startup] {message}", flush=True)
-
-
-def download_model() -> str:
-    MODEL_DIR.mkdir(parents=True, exist_ok=True)
-    local_file = MODEL_DIR / MODEL_FILE
-    if local_file.exists():
-        log(f"Using cached model: {local_file}")
-        return str(local_file)
-
-    log(f"Downloading {MODEL_REPO}/{MODEL_FILE}")
-    model_path = hf_hub_download(
-        repo_id=MODEL_REPO,
-        filename=MODEL_FILE,
-        local_dir=str(MODEL_DIR),
-    )
-    log(f"Model ready: {model_path}")
-    return model_path
-
-
-def download_chat_template() -> str | None:
-    if CHAT_TEMPLATE_FILE.exists() and CHAT_TEMPLATE_FILE.stat().st_size > 0:
-        log(f"Using cached chat template: {CHAT_TEMPLATE_FILE}")
-        return str(CHAT_TEMPLATE_FILE)
-
-    base_repo = "IFM/K2-Horizon-0.9B"
-    encoded_repo = urllib.parse.quote(base_repo, safe="/")
-    api_url = f"https://huggingface.co/api/models/{encoded_repo}"
-    log(f"Fetching chat template from {base_repo} metadata")
-
+import threading
+from typing import Generator
+
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+MODEL_ID = os.getenv("MODEL_ID", "IFM/K2-Horizon-0.9B")
+
+print(f"[startup] Loading tokenizer for {MODEL_ID}...", flush=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+print(f"[startup] Loading model {MODEL_ID}...", flush=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float32,
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+    trust_remote_code=True,
+)
+model.eval()
+print("[startup] Model loaded successfully!", flush=True)
+
+
+def chat_response(
+    message: str,
+    history: list[dict[str, str]],
+    system_prompt: str,
+    temperature: float,
+    top_p: float,
+    max_tokens: int,
+    reasoning_effort: str,
+) -> Generator[str, None, None]:
+    messages = []
+    if system_prompt.strip():
+        messages.append({"role": "system", "content": system_prompt})
+
+    for item in history:
+        messages.append(item)
+
+    messages.append({"role": "user", "content": message})
+
+    chat_kwargs = {"reasoning_effort": reasoning_effort} if reasoning_effort else {}
+    
     try:
-        with urllib.request.urlopen(api_url, timeout=30) as response:
-            metadata = json.loads(response.read().decode("utf-8"))
-    except Exception as exc:
-        log(f"Could not fetch chat template metadata: {exc}")
-        return None
-
-    template = (metadata.get("cardData") or {}).get("chat_template") or metadata.get("chat_template")
-    if not template:
-        log("No chat template found in model metadata; llama-server will use GGUF metadata")
-        return None
-
-    CHAT_TEMPLATE_FILE.parent.mkdir(parents=True, exist_ok=True)
-    CHAT_TEMPLATE_FILE.write_text(template, encoding="utf-8")
-    log(f"Chat template ready: {CHAT_TEMPLATE_FILE}")
-    return str(CHAT_TEMPLATE_FILE)
-
-
-def build_command(model_path: str, template_path: str | None) -> list[str]:
-    def has_custom_value(value: str) -> bool:
-        return value.strip().lower() not in {"", "default", "auto", "none", "off"}
-
-    def add_optional_pair(flag: str, value: str) -> None:
-        if has_custom_value(value):
-            cmd.extend([flag, value])
-
-    cmd = [
-        LLAMA_SERVER_BIN,
-        "-m",
-        model_path,
-        "--host",
-        LLAMA_HOST,
-        "--port",
-        LLAMA_PORT,
-        "--threads",
-        THREADS,
-        "--ctx-size",
-        CTX_SIZE,
-        "--n-gpu-layers",
-        GPU_LAYERS,
-        "--parallel",
-        "1",
-        "--cont-batching",
-        "--temp",
-        TEMPERATURE,
-        "--top-p",
-        TOP_P,
-        "--top-k",
-        TOP_K,
-        "--repeat-penalty",
-        REPEAT_PENALTY,
-    ]
-
-    add_optional_pair("--batch-size", BATCH_SIZE)
-    add_optional_pair("--ubatch-size", UBATCH_SIZE)
-    add_optional_pair("--cache-type-k", CACHE_TYPE_K)
-    add_optional_pair("--cache-type-v", CACHE_TYPE_V)
-    if has_custom_value(FLASH_ATTN):
-        cmd.extend(["-fa", FLASH_ATTN])
-
-    if template_path:
-        cmd.extend(["--chat-template-file", template_path])
-
-    return cmd
-
-
-def main() -> None:
-    binary_dir = str(Path(LLAMA_SERVER_BIN).parent)
-    existing_library_path = os.environ.get("LD_LIBRARY_PATH")
-    os.environ["LD_LIBRARY_PATH"] = (
-        binary_dir if not existing_library_path else f"{binary_dir}:{existing_library_path}"
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            chat_template_kwargs=chat_kwargs if chat_kwargs else None,
+        )
+    except Exception:
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+    inputs = tokenizer([prompt], return_tensors="pt")
+    inputs.pop("token_type_ids", None)
+
+    streamer = TextIteratorStreamer(
+        tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True
     )
 
-    os.environ.setdefault("OMP_NUM_THREADS", THREADS)
-    os.environ.setdefault("OPENBLAS_NUM_THREADS", THREADS)
-    os.environ.setdefault("MKL_NUM_THREADS", THREADS)
-
-    model_path = download_model()
-    template_path = download_chat_template()
-    cmd = build_command(model_path, template_path)
-
-    log("Starting native llama.cpp web UI")
-    log(" ".join(cmd))
-    os.execvpe(cmd[0], cmd, os.environ)
+    generate_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=int(max_tokens),
+        do_sample=temperature > 0,
+        temperature=float(temperature) if temperature > 0 else 1.0,
+        top_p=float(top_p),
+    )
 
+    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
+
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
+        yield partial_text
+
+
+demo = gr.ChatInterface(
+    fn=chat_response,
+    type="messages",
+    title="K2-Horizon-0.9B Chat Demo",
+    description="Interactive demo for [IFM/K2-Horizon-0.9B](https://huggingface.co/IFM/K2-Horizon-0.9B) using PyTorch and Transformers on CPU.",
+    additional_inputs=[
+        gr.Textbox(
+            value="You are a helpful and harmless assistant.",
+            label="System Prompt",
+        ),
+        gr.Slider(
+            minimum=0.0,
+            maximum=2.0,
+            value=0.6,
+            step=0.1,
+            label="Temperature",
+        ),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-P",
+        ),
+        gr.Slider(
+            minimum=128,
+            maximum=8192,
+            value=2048,
+            step=128,
+            label="Max Tokens",
+        ),
+        gr.Dropdown(
+            choices=["high", "medium", "low"],
+            value="high",
+            label="Reasoning Effort",
+        ),
+    ],
+)
 
 if __name__ == "__main__":
-    try:
-        main()
-    except Exception as exc:
-        print(f"[fatal] {exc}", file=sys.stderr, flush=True)
-        raise
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)