File size: 1,057 Bytes
f4c4c18 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | # 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()
|