Use Transformers runtime for Space inference

README.md

@@ -12,7 +12,6 @@ tags:
   - backyard-ai
   - tiny-titan
   - well-tuned
-  - llama-cpp
   - code
 ---
 
@@ -22,7 +21,7 @@ A small-model code-reading assistant for Python files. It parses a single file d
 
 The app includes two modes:
 
-- **Base Mellum2:** richer comments from the base instruct model through llama.cpp.
+- **Base Mellum2:** richer comments from the base instruct model through Transformers.
 - **Fine-tuned LoRA:** a concise comment style trained on CodeSearchNet-derived Python examples with Modal.
 
 The model never edits code directly. It only proposes comments, and the app rejects any annotated file whose semantic AST changes.
@@ -48,7 +47,7 @@ npx between-the-lines-cli path/to/file.py --model base --output annotated.py
 npx between-the-lines-cli path/to/file.py --model tuned --in-place
 ```
 
-`--model base` uses the richer Mellum2 GGUF path. `--model tuned` uses the LoRA adapter for shorter comments. Both modes run the AST validation before writing output.
+`--model base` uses the richer Mellum2 instruct path. `--model tuned` uses the LoRA adapter for shorter comments. Both modes run the AST validation before writing output.
 
 ## Hackathon Fit
 

app.py

@@ -7,6 +7,17 @@ import gradio as gr
 from btl.annotate import annotate_python, collect_blocks, parse_python
 from btl.model import ModelUnavailableError
 
+try:
+    import spaces
+except ImportError:
+    spaces = None
+
+
+def zero_gpu_task(fn):
+    if spaces is None:
+        return fn
+    return spaces.GPU(duration=180)(fn)
+
 
 CSS = """
 :root {
@@ -161,6 +172,7 @@ MODEL_LABELS = {
 }
 
 
+@zero_gpu_task
 def annotate_code(source: str, model_label: str) -> tuple[str, str, str]:
     try:
         result = annotate_python(source, MODEL_LABELS.get(model_label, "base"))

btl/model.py

@@ -5,8 +5,6 @@ from typing import Literal
 from .prompts import build_comment_messages
 
 
-DEFAULT_MODEL_REPO = "JetBrains/Mellum2-12B-A2.5B-Instruct-GGUF-Q8_0"
-DEFAULT_MODEL_FILE = "Mellum2-12B-A2.5B-Instruct-Q8_0.gguf"
 DEFAULT_BASE_TRANSFORMERS_MODEL = "JetBrains/Mellum2-12B-A2.5B-Instruct"
 DEFAULT_TUNED_ADAPTER_REPO = "coolbeanz79/between-the-lines-mellum2-lora"
 
@@ -18,31 +16,38 @@ class ModelUnavailableError(RuntimeError):
 
 
 @lru_cache(maxsize=1)
-def _load_base_llm():
+def _load_base_model():
     try:
-        from llama_cpp import Llama
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
     except ImportError as exc:
         raise ModelUnavailableError(
-            "llama-cpp-python is not installed. Install requirements before using model annotations."
+            "Base Mellum2 inference requires torch, transformers, accelerate, and bitsandbytes."
         ) from exc
 
-    repo_id = os.getenv("BTL_MODEL_REPO", DEFAULT_MODEL_REPO)
-    filename = os.getenv("BTL_MODEL_FILE", DEFAULT_MODEL_FILE)
-    n_ctx = int(os.getenv("BTL_MODEL_CTX", "4096"))
-    n_gpu_layers = int(os.getenv("BTL_MODEL_GPU_LAYERS", "-1"))
+    model_name = os.getenv("BTL_BASE_MODEL", DEFAULT_BASE_TRANSFORMERS_MODEL)
 
     try:
-        return Llama.from_pretrained(
-            repo_id=repo_id,
-            filename=filename,
-            n_ctx=n_ctx,
-            n_gpu_layers=n_gpu_layers,
-            verbose=False,
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
         )
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            device_map="auto",
+            quantization_config=quantization_config,
+        )
+        model.eval()
+        return tokenizer, model
     except Exception as exc:
-        raise ModelUnavailableError(
-            f"Could not load `{repo_id}` / `{filename}` with llama-cpp-python: {exc}"
-        ) from exc
+        raise ModelUnavailableError(f"Could not load base Mellum2 model `{model_name}`: {exc}") from exc
 
 
 def _clean_comment(text: str) -> str:
@@ -74,32 +79,14 @@ def _load_tuned_model():
         )
 
     try:
-        import torch
         from peft import PeftModel
-        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
     except ImportError as exc:
         raise ModelUnavailableError(
-            "Tuned LoRA inference requires torch, transformers, peft, and bitsandbytes."
+            "Tuned LoRA inference requires peft in addition to the base Transformers runtime."
         ) from exc
 
-    model_name = os.getenv("BTL_TUNED_BASE_MODEL", DEFAULT_BASE_TRANSFORMERS_MODEL)
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_use_double_quant=True,
-        )
-        base_model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            device_map="auto",
-            quantization_config=quantization_config,
-        )
+        tokenizer, base_model = _load_base_model()
         model = PeftModel.from_pretrained(base_model, adapter_path_or_repo)
         model.eval()
         return tokenizer, model
@@ -110,39 +97,24 @@ def _load_tuned_model():
 def generate_comment(kind: str, name: str, source: str, variant: ModelVariant = "base") -> str:
     if variant == "tuned":
         return generate_comment_with_tuned_model(kind, name, source)
-    llm = load_llm()
-    return generate_comment_with_llm(llm, kind, name, source)
-
+    return generate_comment_with_base_model(kind, name, source)
 
-def load_llm():
-    return _load_base_llm()
 
-
-def generate_comment_with_llm(llm, kind: str, name: str, source: str) -> str:
-    messages = build_comment_messages(kind, name, source)
-    response = llm.create_chat_completion(
-        messages=messages,
-        temperature=0.15,
-        top_p=0.9,
-        max_tokens=80,
-        stop=["\n\n", "```"],
-    )
-    text = response["choices"][0]["message"]["content"]
-    return _clean_comment(text)
-
-
-def generate_comment_with_tuned_model(kind: str, name: str, source: str) -> str:
+def _generate_with_transformers(tokenizer, model, kind: str, name: str, source: str, max_length_env: str) -> str:
     import torch
 
-    tokenizer, model = _load_tuned_model()
     messages = build_comment_messages(kind, name, source)
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     encoded = tokenizer(
         prompt,
         return_tensors="pt",
         truncation=True,
-        max_length=int(os.getenv("BTL_TUNED_MODEL_CTX", "4096")),
-    ).to(model.device)
+        max_length=int(os.getenv(max_length_env, "4096")),
+    )
+
+    device = getattr(model, "device", None)
+    if device is not None:
+        encoded = encoded.to(device)
 
     with torch.inference_mode():
         output_ids = model.generate(
@@ -156,3 +128,13 @@ def generate_comment_with_tuned_model(kind: str, name: str, source: str) -> str:
     generated_ids = output_ids[encoded["input_ids"].shape[-1] :]
     text = tokenizer.decode(generated_ids, skip_special_tokens=True)
     return _clean_comment(text)
+
+
+def generate_comment_with_base_model(kind: str, name: str, source: str) -> str:
+    tokenizer, model = _load_base_model()
+    return _generate_with_transformers(tokenizer, model, kind, name, source, "BTL_MODEL_CTX")
+
+
+def generate_comment_with_tuned_model(kind: str, name: str, source: str) -> str:
+    tokenizer, model = _load_tuned_model()
+    return _generate_with_transformers(tokenizer, model, kind, name, source, "BTL_TUNED_MODEL_CTX")

pyproject.toml

@@ -13,9 +13,9 @@ dependencies = [
     "bitsandbytes",
     "gradio==5.50.0",
     "huggingface_hub",
-    "llama-cpp-python",
     "peft",
     "safetensors",
+    "spaces",
     "torch",
     "transformers @ git+https://github.com/huggingface/transformers.git",
     "truststore",

requirements.txt

@@ -1,11 +1,10 @@
 gradio==5.50.0
 huggingface_hub
-llama-cpp-python
-datasets==2.21.0
 truststore
 accelerate
 bitsandbytes
 peft
 safetensors
+spaces
 torch
 git+https://github.com/huggingface/transformers.git