backend/providers/gemma.py

"""Swahili Gemma 1B engine via llama-cpp-python (GGUF)."""

from __future__ import annotations

import os
from threading import Lock

from providers.base import TranslationProvider

MODEL_PATH = os.environ.get(
    "GEMMA_MODEL_PATH",
    "/data/models/swahili-gemma-1b/Q4_K_M/swahili-gemma-1b-q4_k_m.gguf",
)


class SwahiliGemmaProvider(TranslationProvider):
    kind = "local"
    description = "Fine-tuned Gemma 1B · Specialized English -> Swahili translation"
    private = True
    setup_hint = "Download GGUF model to /data/models/swahili-gemma-1b"

    def __init__(
        self,
        provider_id: str = "swahili_gemma",
        name: str = "Swahili Gemma (1B GGUF)",
        model_path: str = MODEL_PATH,
    ) -> None:
        self.id = provider_id
        self.name = name
        self.model_path = model_path
        self._llm = None
        self._lock = Lock()

    def is_available(self) -> bool:
        return os.path.isfile(self.model_path)

    def _ensure_loaded(self) -> None:
        if self._llm is not None:
            return
        with self._lock:
            if self._llm is not None:
                return
            from llama_cpp import Llama

            self._llm = Llama(
                model_path=self.model_path,
                n_ctx=2048,
                n_threads=8,  # Match your 8 vCPUs
                verbose=False,
            )

    def translate(self, text: str, src: str, tgt: str, api_key: str | None = None) -> str:
        self._ensure_loaded()
        
        prompt = f"Translate to Swahili: {text}"
        
        response = self._llm(
            prompt,
            max_tokens=256,
            temperature=0.3,
            top_p=0.95,
            repeat_penalty=1.1,
        )
        
        return response["choices"][0]["text"].strip()