Spaces:
Runtime error
Runtime error
File size: 4,331 Bytes
8bc38db | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | 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])
|