inference.py

# inference.py
# ==============================
# Single inference entry point
# Handles Base / Core / Skill transparently
# ==============================

import torch
from model_loader import load_model
from router import route_skill
import config


def generate_response(prompt: str):
    """
    Main inference function.
    Decides skill softly, then generates response.
    """

    # Decide skill (or None)
    skill = route_skill(prompt)

    # Load model (Base + optional Core + optional Skill)
    model, tokenizer = load_model(skill)

    # Prepare input
    inputs = tokenizer(
        prompt,
        return_tensors="pt"
    ).to(model.device)

    # Generate
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=config.MAX_NEW_TOKENS,
            temperature=config.TEMPERATURE,
            top_p=config.TOP_P,
        )

    # Decode
    response = tokenizer.decode(
        output[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True
    )

    return response.strip()