RustLean FIM Space: cursor-aware Rust completion via llama-cpp-python + ZeroGPU

.env.example

@@ -0,0 +1,5 @@
+HF_TOKEN=hf_xxx
+# MODEL_REPO=AfkaraLP/rustlean-gguf
+# MODEL_FILE=rustlean-final.Q8_0.gguf
+# RUSTLEAN_CTX=8192
+# N_GPU_LAYERS=-1

.gitignore

@@ -0,0 +1,7 @@
+.venv/
+.env
+__pycache__/
+*.pyc
+result
+models/
+*.gguf

README.md

@@ -0,0 +1,65 @@
+---
+title: RustLean FIM Playground
+emoji: 🦀
+colorFrom: indigo
+colorTo: red
+sdk: gradio
+sdk_version: "4.44.1"
+app_file: app.py
+pinned: false
+license: apache-2.0
+models:
+  - AfkaraLP/rustlean-gguf
+short_description: Cursor-aware Rust FIM code completion
+---
+
+# RustLean FIM Playground
+
+Cursor-aware Rust code completion using
+[`AfkaraLP/rustlean-gguf`](https://huggingface.co/AfkaraLP/rustlean-gguf) — a
+Qwen2.5-Coder-1.5B model fine-tuned for **fill-in-the-middle (FIM)** on Rust,
+served via `llama-cpp-python`.
+
+## How it works
+
+Place a **`<|cursor|>`** marker anywhere in the editor. The app splits the file
+into a prefix and a suffix at that marker and builds the Qwen PSM prompt:
+
+```
+<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>
+```
+
+The generated middle is spliced back in. If there is no suffix, the model runs
+in pure prefix-completion mode.
+
+## Hardware (ZeroGPU)
+
+Set the Space hardware to **ZeroGPU** in *Settings → Hardware*.
+
+## Secrets
+
+Add a Space **Secret** named `HF_TOKEN` with a read token
+(<https://huggingface.co/settings/tokens>). The Space only needs it if the model
+repo is private or to avoid rate limits while downloading the GGUF.
+
+Optional environment overrides (Space Variables or local `.env`):
+
+| Variable         | Default                    | Meaning                                |
+| ---------------- | -------------------------- | -------------------------------------- |
+| `MODEL_REPO`     | `AfkaraLP/rustlean-gguf`   | HF repo id                             |
+| `MODEL_FILE`     | `rustlean-final.Q8_0.gguf` | GGUF filename                          |
+| `RUSTLEAN_CTX`   | `8192`                     | Context window                         |
+| `N_GPU_LAYERS`   | `-1`                       | `-1` offloads all layers to the GPU    |
+
+## Local development (NixOS / CUDA)
+
+This repo ships a `flake.nix` for a CUDA dev shell (built for the RTX 2070 Super,
+compute capability `sm_75` / Turing):
+
+```bash
+nix develop         # sets up CUDA + builds llama-cpp-python with GGML_CUDA
+cp .env.example .env && edit .env   # put your HF_TOKEN here
+python app.py
+```
+
+On a 2070 Super expect ~40 tok/s decode (per the model card).

app.py

@@ -0,0 +1,209 @@
+import os
+import time
+from pathlib import Path
+
+
+def _load_env(path: str = ".env"):
+    p = Path(path)
+    if not p.exists():
+        return
+    for line in p.read_text().splitlines():
+        line = line.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+        k, v = line.split("=", 1)
+        os.environ.setdefault(k.strip(), v.strip())
+
+
+_load_env()
+
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
+
+try:
+    import spaces
+    gpu = spaces.GPU
+except Exception:
+    def gpu(*args, **kwargs):
+        def deco(fn):
+            return fn
+        return deco
+
+MODEL_REPO = os.environ.get("MODEL_REPO", "AfkaraLP/rustlean-gguf")
+MODEL_FILE = os.environ.get("MODEL_FILE", "rustlean-final.Q8_0.gguf")
+CTX = int(os.environ.get("RUSTLEAN_CTX", "8192"))
+N_GPU_LAYERS = int(os.environ.get("N_GPU_LAYERS", "-1"))
+
+FIM_PREFIX = "<|fim_prefix|>"
+FIM_SUFFIX = "<|fim_suffix|>"
+FIM_MIDDLE = "<|fim_middle|>"
+CURSOR = "<|cursor|>"
+STOPS = ["<|endoftext|>", FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE, "<|fim_pad|>"]
+
+DEFAULT_CODE = """use std::collections::HashMap;
+
+fn word_count(text: &str) -> HashMap<&str, u32> {
+    let mut counts = HashMap::new();
+<|cursor|>
+    counts
+}
+"""
+
+EXAMPLES = [
+    ("""fn fibonacci(n: u32) -> u32 {
+    match n {
+        0 => 0,
+        1 => 1,
+<|cursor|>
+    }
+}
+""", 0.2, 96, 0.9),
+    ("""struct Point { x: f64, y: f64 }
+
+impl Point {
+    fn distance(&self, other: &Point) -> f64 {
+<|cursor|>
+    }
+}
+""", 0.2, 96, 0.9),
+    ("""fn is_prime(n: u64) -> bool {
+    if n < 2 {
+        return false;
+    }
+<|cursor|>
+}
+""", 0.2, 128, 0.9),
+]
+
+_llm = None
+_model_path = None
+
+
+def model_path():
+    global _model_path
+    if _model_path is None:
+        _model_path = hf_hub_download(
+            repo_id=MODEL_REPO,
+            filename=MODEL_FILE,
+            token=os.environ.get("HF_TOKEN"),
+        )
+    return _model_path
+
+
+def get_llm():
+    global _llm
+    if _llm is None:
+        _llm = Llama(
+            model_path=model_path(),
+            n_gpu_layers=N_GPU_LAYERS,
+            n_ctx=CTX,
+            n_batch=512,
+            verbose=False,
+            logits_all=False,
+        )
+    return _llm
+
+
+def build_prompt(prefix: str, suffix: str) -> str:
+    return f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}"
+
+
+def clean(text: str) -> str:
+    for s in [FIM_MIDDLE, FIM_SUFFIX, FIM_PREFIX, "<|endoftext|>", "<|fim_pad|>"]:
+        text = text.replace(s, "")
+    return text.rstrip()
+
+
+@gpu(duration=60)
+def complete(code, temperature, max_tokens, top_p):
+    if code.strip() == "":
+        return "", code, "*Empty input.*"
+    if CURSOR in code:
+        prefix, suffix = code.split(CURSOR, 1)
+    else:
+        prefix, suffix = code, ""
+    prompt = build_prompt(prefix, suffix)
+
+    def run():
+        llm = get_llm()
+        return llm.create_completion(
+            prompt=prompt,
+            max_tokens=int(max_tokens),
+            temperature=float(temperature),
+            top_p=float(top_p),
+            stop=STOPS,
+            stream=False,
+        )
+
+    t0 = time.time()
+    try:
+        out = run()
+    except Exception:
+        global _llm
+        _llm = None
+        out = run()
+    dt = time.time() - t0
+
+    text = clean(out["choices"][0]["text"])
+    assembled = f"{prefix}{text}{suffix}"
+    n_tok = out.get("usage", {}).get("completion_tokens", 0)
+    tps = (n_tok / dt) if dt > 0 else 0.0
+    mode = "prefix" if suffix == "" else "FIM (prefix+suffix)"
+    stats = (
+        f"**{n_tok}** tokens · **{dt:.2f}s** · **{tps:.1f} tok/s**"
+        f"  ·  mode: {mode}"
+        f"  ·  gpu_layers: {N_GPU_LAYERS}"
+    )
+    return text, assembled, stats
+
+
+with gr.Blocks(title="RustLean FIM Playground", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        "# 🦀 RustLean — Rust FIM Completion\n"
+        "Fill-in-the-middle completion from "
+        "[`AfkaraLP/rustlean-gguf`](https://huggingface.co/AfkaraLP/rustlean-gguf) "
+        "(Qwen2.5-Coder-1.5B, Q8_0). Put **`<|cursor|>`** where you want the model to "
+        "write code. Everything before it is the prefix, everything after is the suffix."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            code_in = gr.Code(
+                value=DEFAULT_CODE,
+                language="rust",
+                label="Rust source (mark the hole with <|cursor|>)",
+                interactive=True,
+            )
+            with gr.Accordion("Sampling", open=False):
+                temperature = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="temperature")
+                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
+                max_tokens = gr.Slider(16, 256, value=96, step=8, label="max tokens")
+            with gr.Row():
+                btn = gr.Button("Complete at cursor", variant="primary")
+                clear = gr.Button("Reset")
+        with gr.Column(scale=1):
+            stats = gr.Markdown("_Not run yet._")
+            completion = gr.Code(
+                label="Generated middle",
+                language="rust",
+                interactive=False,
+            )
+            result = gr.Code(
+                label="Assembled result (prefix + middle + suffix)",
+                language="rust",
+                interactive=True,
+            )
+
+    gr.Examples(examples=EXAMPLES, inputs=[code_in, temperature, max_tokens, top_p])
+
+    btn.click(
+        complete,
+        inputs=[code_in, temperature, max_tokens, top_p],
+        outputs=[completion, result, stats],
+    )
+    clear.click(lambda: (DEFAULT_CODE, "", "", "_Reset._"), None,
+                outputs=[code_in, completion, result, stats])
+
+if __name__ == "__main__":
+    demo.queue().launch()

flake.nix

@@ -0,0 +1,58 @@
+{
+  description = "RustLean FIM Space — local dev shell (CUDA + llama-cpp-python) for RTX 2070 Super";
+
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+    flake-utils.url = "github:numtide/flake-utils";
+  };
+
+  outputs = { self, nixpkgs, flake-utils }:
+    flake-utils.lib.eachDefaultSystem (system:
+      let
+        pkgs = import nixpkgs {
+          inherit system;
+          config.allowUnfree = true;
+        };
+        cuda = pkgs.cudaPackages;
+        py = pkgs.python312;
+      in {
+        devShells.default = pkgs.mkShell {
+          packages = [
+            py
+            py.pkgs.pip
+            py.pkgs.virtualenv
+            pkgs.cmake
+            pkgs.ninja
+            pkgs.gcc
+            pkgs.gnumake
+            pkgs.git
+            cuda.cudatoolkit
+            cuda.cudnn
+          ];
+
+          env = {
+            CUDA_HOME = "${cuda.cudatoolkit}";
+            CMAKE_ARGS = "-DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=75";
+          };
+
+          shellHook = ''
+            export PATH="$CUDA_HOME/bin:$PATH"
+            export LD_LIBRARY_PATH="${cuda.cudatoolkit}/lib:${cuda.cudnn}/lib:$LD_LIBRARY_PATH"
+            export LIBRARY_PATH="${cuda.cudatoolkit}/lib:${cuda.cudnn}/lib:$LIBRARY_PATH"
+
+            if [ ! -d .venv ]; then
+              ${py}/bin/python -m venv .venv
+              ./.venv/bin/pip install --upgrade pip wheel setuptools
+            fi
+            source .venv/bin/activate
+
+            pip install -r requirements.txt
+            pip install --force-reinstall --no-binary llama-cpp-python llama-cpp-python
+
+            echo ""
+            echo "RustLean dev shell ready.  RTX 2070 Super = sm_75 (Turing)."
+            echo "Put your token in .env (HF_TOKEN), then run:  python app.py"
+          '';
+        };
+      });
+}

requirements.txt

@@ -0,0 +1,4 @@
+--extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu124
+gradio>=4.44.0
+huggingface_hub>=0.25.0
+llama-cpp-python>=0.3.1