src/inference_engine.py

import numpy as np
import onnxruntime as ort
from typing import Any


class InferenceEngine:
    def __init__(self, session: ort.InferenceSession, tokenizer: Any, config: Any):
        self.session = session
        self.tokenizer = tokenizer
        self.config = config
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name

    def generate_response(self, prompt: str, max_tokens: int = 64, temperature: float = 0.8,
                          top_p: float = 0.9, top_k: int = 50) -> str:
        # Encode prompt
        tokens = self.tokenizer.encode(prompt)
        input_ids = np.array([tokens], dtype=np.int64)
        generated = []

        for _ in range(max_tokens):
            outputs = self.session.run([self.output_name], {self.input_name: input_ids})
            logits = outputs[0][0, -1, :]  # last token logits

            # Temperature
            if temperature and temperature > 0:
                logits = logits / max(temperature, 1e-6)

            # Top-k
            if top_k and top_k > 0:
                k = min(top_k, logits.shape[-1])
                idx = np.argpartition(logits, -k)[-k:]
                filt = np.full_like(logits, -np.inf)
                filt[idx] = logits[idx]
                logits = filt

            # Softmax
            exps = np.exp(logits - np.max(logits))
            probs = exps / np.sum(exps)

            # Top-p (nucleus)
            if top_p is not None and 0 < top_p < 1.0:
                sort_idx = np.argsort(probs)[::-1]
                sorted_probs = probs[sort_idx]
                cumsum = np.cumsum(sorted_probs)
                cutoff = np.searchsorted(cumsum, top_p) + 1
                mask = np.zeros_like(probs)
                keep = sort_idx[:cutoff]
                mask[keep] = probs[keep]
                s = mask.sum()
                if s > 0:
                    probs = mask / s

            next_token = int(np.random.choice(len(probs), p=probs))
            if next_token == self.tokenizer.eos_token_id:
                break

            generated.append(next_token)
            input_ids = np.concatenate([input_ids, [[next_token]]], axis=1)

        text = self.tokenizer.decode(generated, skip_special_tokens=True).strip()
        if not text:
            return "I couldn't generate a response."
        # Remove prompt echo if present (best-effort)
        if text.startswith(prompt):
            text = text[len(prompt):].strip()
        return text