libs

libs/inference/README.md

@@ -0,0 +1,11 @@
+# inference
+
+Swappable local inference backends (`llama_cpp` default, `transformers` optional extra).
+
+```python
+from inference.factory import get_backend
+
+backend = get_backend()
+backend.load()
+reply = backend.chat([{"role": "user", "content": "Hello!"}])
+```

libs/inference/pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "inference"
+version = "0.1.0"
+description = "Swappable local inference backends for the hackathon Space"
+readme = "README.md"
+authors = [
+    { name = "MSGhais", email = "msghais135@gmail.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "huggingface-hub>=0.27.0",
+    "llama-cpp-python>=0.3.0",
+]
+
+[project.optional-dependencies]
+transformers = [
+    "accelerate>=1.2.0",
+    "torch>=2.5.0",
+    "transformers>=4.47.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.8.13,<0.9.0"]
+build-backend = "uv_build"

libs/inference/src/inference/__init__.py

@@ -0,0 +1,3 @@
+from inference.factory import get_backend
+
+__all__ = ["get_backend"]

libs/inference/src/inference/base.py

@@ -0,0 +1,24 @@
+from typing import Protocol
+
+
+class InferenceBackend(Protocol):
+    def load(self) -> None:
+        """Load model weights into memory."""
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ) -> str:
+        """Generate text from a single prompt."""
+
+    def chat(
+        self,
+        messages: list[dict[str, str]],
+        *,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ) -> str:
+        """Generate a reply from a chat message history."""

libs/inference/src/inference/factory.py

@@ -0,0 +1,23 @@
+import os
+from functools import lru_cache
+
+from inference.base import InferenceBackend
+from inference.llama_cpp import LlamaCppBackend
+
+
+@lru_cache(maxsize=1)
+def get_backend() -> InferenceBackend:
+    backend_name = os.environ.get("INFERENCE_BACKEND", "llama_cpp").lower()
+
+    if backend_name == "llama_cpp":
+        return LlamaCppBackend()
+
+    if backend_name == "transformers":
+        from inference.transformers import TransformersBackend
+
+        return TransformersBackend()
+
+    raise ValueError(
+        f"Unknown INFERENCE_BACKEND={backend_name!r}. "
+        "Expected 'llama_cpp' or 'transformers'."
+    )

libs/inference/src/inference/llama_cpp.py

@@ -0,0 +1,83 @@
+import os
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
+
+
+DEFAULT_MODEL_REPO = "Qwen/Qwen2.5-3B-Instruct-GGUF"
+DEFAULT_MODEL_FILE = "qwen2.5-3b-instruct-q4_k_m.gguf"
+
+
+class LlamaCppBackend:
+    def __init__(self) -> None:
+        self._model: Llama | None = None
+        self._model_path: str | None = None
+
+    def _resolve_model_path(self) -> str:
+        model_path = os.environ.get("MODEL_PATH")
+        if model_path:
+            path = Path(model_path)
+            if not path.exists():
+                raise FileNotFoundError(f"MODEL_PATH does not exist: {model_path}")
+            return str(path)
+
+        model_repo = os.environ.get("MODEL_REPO", DEFAULT_MODEL_REPO)
+        model_file = os.environ.get("MODEL_FILE", DEFAULT_MODEL_FILE)
+        cache_dir = os.environ.get("MODEL_CACHE_DIR")
+
+        return hf_hub_download(
+            repo_id=model_repo,
+            filename=model_file,
+            cache_dir=cache_dir,
+        )
+
+    def load(self) -> None:
+        if self._model is not None:
+            return
+
+        self._model_path = self._resolve_model_path()
+        n_ctx = int(os.environ.get("N_CTX", "4096"))
+        n_gpu_layers = int(os.environ.get("N_GPU_LAYERS", "0"))
+
+        self._model = Llama(
+            model_path=self._model_path,
+            n_ctx=n_ctx,
+            n_gpu_layers=n_gpu_layers,
+            verbose=False,
+        )
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ) -> str:
+        self.load()
+        assert self._model is not None
+
+        result = self._model(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            echo=False,
+        )
+        return result["choices"][0]["text"].strip()
+
+    def chat(
+        self,
+        messages: list[dict[str, str]],
+        *,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ) -> str:
+        self.load()
+        assert self._model is not None
+
+        result = self._model.create_chat_completion(
+            messages=messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+        return result["choices"][0]["message"]["content"].strip()

libs/inference/src/inference/transformers.py

@@ -0,0 +1,89 @@
+import os
+
+from inference.base import InferenceBackend
+
+
+class TransformersBackend:
+    def __init__(self) -> None:
+        self._model = None
+        self._tokenizer = None
+
+    def load(self) -> None:
+        if self._model is not None:
+            return
+
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as exc:
+            raise ImportError(
+                "transformers backend requires optional deps. "
+                "Install with: uv sync --package inference --extra transformers"
+            ) from exc
+
+        model_id = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-3B-Instruct")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self._tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self._model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            device_map="auto" if device == "cuda" else None,
+        )
+        if device == "cpu":
+            self._model.to(device)
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ) -> str:
+        self.load()
+        assert self._model is not None
+        assert self._tokenizer is not None
+
+        import torch
+
+        inputs = self._tokenizer(prompt, return_tensors="pt").to(self._model.device)
+        output = self._model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            do_sample=temperature > 0,
+        )
+        generated = output[0][inputs["input_ids"].shape[-1] :]
+        return self._tokenizer.decode(generated, skip_special_tokens=True).strip()
+
+    def chat(
+        self,
+        messages: list[dict[str, str]],
+        *,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ) -> str:
+        self.load()
+        assert self._model is not None
+        assert self._tokenizer is not None
+
+        if hasattr(self._tokenizer, "apply_chat_template"):
+            prompt = self._tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+        else:
+            parts = []
+            for message in messages:
+                role = message["role"]
+                content = message["content"]
+                parts.append(f"{role}: {content}")
+            parts.append("assistant:")
+            prompt = "\n".join(parts)
+
+        return self.generate(prompt, max_tokens=max_tokens, temperature=temperature)
+
+
+# Satisfy static type checkers that expect InferenceBackend.
+_: InferenceBackend = TransformersBackend()