topicjev.backend

Shared device handling and model lifecycle. Every backend in TopicJev derives from ModelBackend. See Hardware and Memory.

Devices

detect_device

detect_device() -> tuple[device, dtype]

Select the optimal torch device and dtype for local inference.

Prefers CUDA, then Apple Silicon (MPS), then CPU. Set the TOPICJEV_DEVICE environment variable (e.g. cpu, mps, cuda:1) to override.

Source code in topicjev/backend.py

def detect_device() -> tuple[torch.device, torch.dtype]:
    """Select the optimal torch device and dtype for local inference.

    Prefers CUDA, then Apple Silicon (MPS), then CPU. Set the ``TOPICJEV_DEVICE``
    environment variable (e.g. ``cpu``, ``mps``, ``cuda:1``) to override.
    """
    override = os.environ.get("TOPICJEV_DEVICE")
    if override:
        dev = torch.device(override)
    elif torch.cuda.is_available():
        dev = torch.device("cuda")
    elif torch.backends.mps.is_available():
        dev = torch.device("mps")
    else:
        dev = torch.device("cpu")

    if dev.type == "cuda":
        major, _ = torch.cuda.get_device_capability(dev)
        dtype = torch.bfloat16 if major >= 8 else torch.float16
    else:
        # MPS stays in float32 so scores match CPU runs
        dtype = torch.float32
    return dev, dtype

empty_device_cache

empty_device_cache() -> None

Release cached accelerator memory (CUDA or MPS) after a model is offloaded.

Source code in topicjev/backend.py

def empty_device_cache() -> None:
    """Release cached accelerator memory (CUDA or MPS) after a model is offloaded."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    elif torch.backends.mps.is_available():
        torch.mps.empty_cache()

Lifecycle

ModelBackend

ModelBackend(mname: str | None = None, **kwargs: Any)

Bases: ABC

Base contract for model and API client lifecycles.

Source code in topicjev/backend.py

def __init__(self, mname: str | None = None, **kwargs: Any) -> None:
    self.mname = mname
    self.model: Any = None

load_model abstractmethod

load_model() -> None

Load the model or initialize the client.

Source code in topicjev/backend.py

@abstractmethod
def load_model(self) -> None:
    """Load the model or initialize the client."""

close

close() -> None

Release allocated model resources.

Source code in topicjev/backend.py

def close(self) -> None:
    """Release allocated model resources."""
    self.model = None

LocalModelBackend

LocalModelBackend(mname: str | None = None, **kwargs: Any)

Bases: ModelBackend

Base for local PyTorch-based models with device and memory management.

Source code in topicjev/backend.py

def __init__(self, mname: str | None = None, **kwargs: Any) -> None:
    super().__init__(mname, **kwargs)
    self.dev, self.dtype = detect_device()
    self.tok: Any = None

close

close() -> None

Offload model to CPU and clear GPU cache if allocated.

Source code in topicjev/backend.py

def close(self) -> None:
    """Offload model to CPU and clear GPU cache if allocated."""
    if self.model is not None and hasattr(self.model, "cpu"):
        try:
            self.model.cpu()
        except Exception as e:
            print("Failed to move model to cpu during close: %s", e)
    super().close()
    self.tok = None
    empty_device_cache()

CausalLMBackend

CausalLMBackend(mname: str | None = None, **kwargs: Any)

Bases: LocalModelBackend, ChatTemplateMixin

Base for local causal language models with left-padded tokenizers.

Source code in topicjev/backend.py

def __init__(self, mname: str | None = None, **kwargs: Any) -> None:
    super().__init__(mname, **kwargs)
    self.dev, self.dtype = detect_device()
    self.tok: Any = None

load_model

load_model() -> None

Load causal LM and left-padded tokenizer onto target hardware device.

Source code in topicjev/backend.py

def load_model(self) -> None:
    """Load causal LM and left-padded tokenizer onto target hardware device."""
    from transformers import AutoModelForCausalLM, AutoTokenizer

    self.model = AutoModelForCausalLM.from_pretrained(
        self.mname,
        dtype=self.dtype,
        attn_implementation="sdpa",
    ).to(self.dev)
    self.tok = AutoTokenizer.from_pretrained(self.mname)
    self.tok.padding_side = "left"
    if self.tok.pad_token is None:
        self.tok.pad_token = self.tok.eos_token
    self.model.eval()

Prompt formatting

ChatTemplateMixin

Mixin providing chat-template prompt formatting for causal / chat models.

format_chat_prompt

format_chat_prompt(
    tokenizer: Any,
    text: str,
    sys_prompt: str | None = None,
    *,
    chat: bool = True,
    thinking: bool = False,
    add_generation_prompt: bool = True,
) -> str

Format input text using the tokenizer's chat template if available.

Parameters:

Name	Type	Description	Default
tokenizer	Any	Hugging Face tokenizer instance.	required
text	str	User query or prompt text.	required
sys_prompt	str | None	Optional system prompt to prepend.	None
chat	bool	Whether to apply chat templating (if supported by tokenizer).	True
thinking	bool	Whether to keep thinking enabled for reasoning-tuned models.	False
add_generation_prompt	bool	Whether to append the generation prompt.	True

Returns:

Type	Description
str	Formatted prompt string.

Source code in topicjev/backend.py

def format_chat_prompt(
    tokenizer: Any,
    text: str,
    sys_prompt: str | None = None,
    *,
    chat: bool = True,
    thinking: bool = False,
    add_generation_prompt: bool = True,
) -> str:
    """Format input text using the tokenizer's chat template if available.

    Args:
        tokenizer: Hugging Face tokenizer instance.
        text: User query or prompt text.
        sys_prompt: Optional system prompt to prepend.
        chat: Whether to apply chat templating (if supported by tokenizer).
        thinking: Whether to keep thinking enabled for reasoning-tuned models.
        add_generation_prompt: Whether to append the generation prompt.

    Returns:
        Formatted prompt string.
    """
    if not chat or not tokenizer or not getattr(tokenizer, "chat_template", None):
        return text

    messages: list[dict[str, str]] = []
    if sys_prompt:
        messages.append({"role": "system", "content": sys_prompt})
    messages.append({"role": "user", "content": text})

    kwargs: dict[str, Any] = {
        "tokenize": False,
        "add_generation_prompt": add_generation_prompt,
    }
    if not thinking:
        kwargs["enable_thinking"] = False

    try:
        return tokenizer.apply_chat_template(messages, **kwargs)
    except TypeError:
        # Checkpoint chat template does not accept enable_thinking
        kwargs.pop("enable_thinking", None)
        return tokenizer.apply_chat_template(messages, **kwargs)

resolve_max_context_length

resolve_max_context_length(
    tokenizer: Any, model: Any = None, default: int = 512
) -> int

Determine the effective maximum sequence length from tokenizer or model config.

Source code in topicjev/backend.py

def resolve_max_context_length(
    tokenizer: Any,
    model: Any = None,
    default: int = 512,
) -> int:
    """Determine the effective maximum sequence length from tokenizer or model config."""
    n = getattr(tokenizer, "model_max_length", None)
    if n is None or n > 1_000_000:
        cfg = getattr(model, "config", None)
        n = getattr(
            cfg,
            "n_positions",
            getattr(cfg, "max_position_embeddings", default),
        )
    return int(n if n is not None else default)