Spaces:
Runtime error
Runtime error
| import os | |
| import sys | |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) | |
| import torch | |
| from ml.model import TheAbyssModel, AbyssConfig | |
| from ml.tokenizer.tokenizer import AbyssTokenizer | |
| from ml.tokenizer.special_tokens import DEFAULT_SYSTEM | |
| _model = None | |
| _tok = None | |
| _device = None | |
| _cfg = None | |
| def load_model(checkpoint_path: str) -> None: | |
| global _model, _tok, _device, _cfg | |
| _device = "cuda" if torch.cuda.is_available() else "cpu" | |
| ckpt = torch.load(checkpoint_path, map_location=_device, weights_only=False) | |
| cfg_raw = ckpt["config"] | |
| _cfg = AbyssConfig(**cfg_raw) if isinstance(cfg_raw, dict) else cfg_raw | |
| _model = TheAbyssModel(_cfg) | |
| _model.load_state_dict(ckpt["model"]) | |
| _model.to(_device) | |
| _model.eval() | |
| tok_path = os.path.join(os.path.dirname(__file__), "..", "..", "ml", "tokenizer", "vocab") | |
| _tok = AbyssTokenizer.load(tok_path) | |
| print(f"[inference] model loaded on {_device}") | |
| def _build_prompt(messages: list) -> str: | |
| # The model was SFT'd with a system turn on EVERY example. If the client sent | |
| # no system message, inject the default one so the served prompt matches the | |
| # training distribution (otherwise the model drifts off the user's request). | |
| if not any(m.get("role") == "system" for m in messages): | |
| messages = [{"role": "system", "content": DEFAULT_SYSTEM}] + list(messages) | |
| parts = [] | |
| for m in messages: | |
| role = m["role"] | |
| content = m["content"] | |
| if role == "system": | |
| parts.append(f"<|system|>{content}<|endofturn|>") | |
| elif role == "user": | |
| parts.append(f"<|user|>{content}<|endofturn|>") | |
| else: | |
| parts.append(f"<|assistant|>{content}<|endofturn|>") | |
| parts.append("<|assistant|>") | |
| return "".join(parts) | |
| def generate_tokens( | |
| messages: list, | |
| max_tokens: int = 200, | |
| temperature: float = 0.7, | |
| top_k: int = 40, | |
| top_p: float = 0.95, | |
| repetition_penalty: float = 1.2, | |
| ): | |
| """Sync generator - yields decoded token strings one at a time.""" | |
| prompt = _build_prompt(messages) | |
| ids = _tok.encode(prompt) | |
| block_size = _cfg.block_size | |
| # Never ask for more output than the context can hold, and always leave | |
| # room for the prompt so the input can't collapse to an empty sequence. | |
| max_tokens = max(1, min(max_tokens, block_size - 8)) | |
| keep = block_size - max_tokens | |
| if len(ids) > keep: | |
| ids = ids[-keep:] | |
| idx = torch.tensor([ids], dtype=torch.long, device=_device) | |
| eot_ids = _tok.encode("<|endofturn|>") | |
| eos_ids = _tok.encode("<|endoftext|>") | |
| stop_ids = set() | |
| if eot_ids: | |
| stop_ids.add(eot_ids[0]) | |
| if eos_ids: | |
| stop_ids.add(eos_ids[0]) | |
| generated: list[int] = [] # tokens produced so far (for repetition penalty) | |
| with torch.no_grad(): | |
| for _ in range(max_tokens): | |
| logits, _ = _model(idx[:, -block_size:]) | |
| logits = logits[0, -1, :] / max(temperature, 1e-8) | |
| # Repetition penalty (CTRL-style) over already-generated tokens only, | |
| # so we damp loops without suppressing words from the user's prompt. | |
| if repetition_penalty and repetition_penalty != 1.0 and generated: | |
| for tid in set(generated): | |
| if logits[tid] > 0: | |
| logits[tid] /= repetition_penalty | |
| else: | |
| logits[tid] *= repetition_penalty | |
| if top_k > 0: | |
| k = min(top_k, logits.size(-1)) | |
| topk_vals, _ = torch.topk(logits, k) | |
| logits = logits.masked_fill(logits < topk_vals[-1], float("-inf")) | |
| probs = torch.softmax(logits, dim=-1) | |
| sorted_probs, sorted_idx = torch.sort(probs, descending=True) | |
| cum = torch.cumsum(sorted_probs, dim=-1) | |
| remove = (cum - sorted_probs) > top_p | |
| sorted_probs[remove] = 0.0 | |
| sorted_probs = sorted_probs / sorted_probs.sum() | |
| next_rel = torch.multinomial(sorted_probs, 1) | |
| next_tok = sorted_idx[next_rel] | |
| tok_id = next_tok.item() | |
| if tok_id in stop_ids: | |
| break | |
| generated.append(tok_id) | |
| idx = torch.cat([idx, next_tok.view(1, 1)], dim=1) | |
| yield _tok.decode([tok_id]) | |