File size: 8,694 Bytes
a006332 | 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | """Text generation and chat helpers for Delta Ultra Mini."""
from __future__ import annotations
import logging
import os
from pathlib import Path
from typing import Any, Generator
import torch
from torch.nn import functional as F
from delta.identity import identity_response
from delta.model import DeltaConfig, DeltaModel
from delta.tokenizer import DEFAULT_SYSTEM_PROMPT, DeltaTokenizer
logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)
class DeltaGenerator:
"""Autoregressive generator for Delta Ultra Mini."""
def __init__(self, model: DeltaModel, tokenizer: DeltaTokenizer, device: str | torch.device | None = None) -> None:
self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
self.model = model.to(self.device)
self.model.eval()
self.tokenizer = tokenizer
@classmethod
def from_files(
cls,
checkpoint_path: str | Path,
tokenizer_path: str | Path,
config_path: str | Path | None = None,
device: str | torch.device | None = None,
) -> "DeltaGenerator":
"""Create a generator from checkpoint, tokenizer, and optional config."""
checkpoint = torch.load(checkpoint_path, map_location="cpu")
config_data = checkpoint.get("config")
config = DeltaConfig.from_dict(config_data) if config_data else DeltaConfig.from_json(config_path or "configs/ultra_mini.json")
model = DeltaModel(config)
state = checkpoint.get("model_state_dict", checkpoint)
model.load_state_dict(state)
return cls(model=model, tokenizer=DeltaTokenizer(tokenizer_path), device=device)
def _filter_logits(
self,
logits: torch.Tensor,
temperature: float,
top_k: int,
top_p: float,
) -> torch.Tensor:
"""Apply temperature, top-k, and nucleus filtering."""
if temperature <= 0:
return logits
logits = logits / temperature
if top_k > 0:
values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits = logits.masked_fill(logits < values[:, [-1]], torch.finfo(logits.dtype).min)
if 0.0 < top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
sorted_probs = F.softmax(sorted_logits, dim=-1)
cumulative = sorted_probs.cumsum(dim=-1)
remove = cumulative > top_p
remove[..., 1:] = remove[..., :-1].clone()
remove[..., 0] = False
indices_to_remove = remove.scatter(1, sorted_indices, remove)
logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
return logits
def _apply_repetition_penalty(
self,
logits: torch.Tensor,
generated: torch.Tensor,
repetition_penalty: float,
) -> torch.Tensor:
"""Penalize tokens that already appeared in the generated sequence."""
if repetition_penalty == 1.0:
return logits
for token_id in set(generated[0].tolist()):
score = logits[:, token_id]
logits[:, token_id] = torch.where(score < 0, score * repetition_penalty, score / repetition_penalty)
return logits
def _clean_completion_text(self, text: str) -> str:
"""Trim leaked prompt/chat markers from decoded assistant text."""
cut_markers = ("[SYS]", "[USR]", "[ASS]", "[SEP]", DEFAULT_SYSTEM_PROMPT)
clean = text
for marker in cut_markers:
index = clean.find(marker)
if index >= 0:
clean = clean[:index]
return clean.strip()
@torch.inference_mode()
def generate(
self,
input_ids: list[int] | torch.Tensor,
max_new_tokens: int = 256,
temperature: float = 0.2,
top_k: int = 20,
top_p: float = 0.9,
repetition_penalty: float = 1.08,
) -> list[int]:
"""Generate token ids using KV cache and manual sampling."""
ids = torch.tensor([input_ids], dtype=torch.long, device=self.device) if isinstance(input_ids, list) else input_ids.to(self.device)
if ids.dim() == 1:
ids = ids.unsqueeze(0)
ids = ids[:, -(self.model.config.max_seq_len - 1) :]
max_new_tokens = min(max_new_tokens, self.model.config.max_seq_len - ids.size(1))
generated = ids.clone()
past_key_values = None
next_input = ids
stop_token_ids = self.tokenizer.chat_stop_token_ids
for _ in range(max_new_tokens):
outputs = self.model(next_input, past_key_values=past_key_values, use_cache=True)
logits = outputs["logits"][:, -1, :]
past_key_values = outputs["past_key_values"]
logits = self._apply_repetition_penalty(logits, generated, repetition_penalty)
if temperature <= 0:
next_token = torch.argmax(logits, dim=-1, keepdim=True)
else:
filtered = self._filter_logits(logits, temperature=temperature, top_k=top_k, top_p=top_p)
next_token = torch.multinomial(F.softmax(filtered, dim=-1), num_samples=1)
if int(next_token.item()) in stop_token_ids:
break
generated = torch.cat((generated, next_token), dim=1)
next_input = next_token
return generated[0].tolist()
def chat(self, messages: list[dict[str, Any]], persona: str | None = None, **gen_kwargs: Any) -> str:
"""Generate an assistant response for chat messages."""
latest_user = next((str(m.get("content", "")) for m in reversed(messages) if m.get("role") == "user"), "")
intercepted = identity_response(latest_user)
if intercepted is not None:
return intercepted
prompt = self.tokenizer.format_chat(messages, persona=persona)
input_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
input_ids = input_ids[-(self.model.config.max_seq_len - 1) :]
output_ids = self.generate(input_ids, **gen_kwargs)
new_ids = output_ids[len(input_ids) :]
text = self.tokenizer.decode(new_ids, skip_special_tokens=True)
return self._clean_completion_text(text)
def stream_chat(
self,
messages: list[dict[str, Any]],
persona: str | None = None,
**gen_kwargs: Any,
) -> Generator[str, None, None]:
"""Yield generated text token by token."""
latest_user = next((str(m.get("content", "")) for m in reversed(messages) if m.get("role") == "user"), "")
intercepted = identity_response(latest_user)
if intercepted is not None:
yield intercepted
return
prompt = self.tokenizer.format_chat(messages, persona=persona)
input_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
input_ids = input_ids[-(self.model.config.max_seq_len - 1) :]
ids = torch.tensor([input_ids], dtype=torch.long, device=self.device)
generated = ids.clone()
past_key_values = None
next_input = ids
max_new_tokens = int(gen_kwargs.get("max_new_tokens", 256))
max_new_tokens = min(max_new_tokens, self.model.config.max_seq_len - ids.size(1))
temperature = float(gen_kwargs.get("temperature", 0.2))
top_k = int(gen_kwargs.get("top_k", 20))
top_p = float(gen_kwargs.get("top_p", 0.9))
repetition_penalty = float(gen_kwargs.get("repetition_penalty", 1.08))
stop_token_ids = self.tokenizer.chat_stop_token_ids
with torch.inference_mode():
for _ in range(max_new_tokens):
outputs = self.model(next_input, past_key_values=past_key_values, use_cache=True)
logits = outputs["logits"][:, -1, :]
past_key_values = outputs["past_key_values"]
logits = self._apply_repetition_penalty(logits, generated, repetition_penalty)
if temperature <= 0:
next_token = torch.argmax(logits, dim=-1, keepdim=True)
else:
filtered = self._filter_logits(logits, temperature=temperature, top_k=top_k, top_p=top_p)
next_token = torch.multinomial(F.softmax(filtered, dim=-1), num_samples=1)
if int(next_token.item()) in stop_token_ids:
return
generated = torch.cat((generated, next_token), dim=1)
next_input = next_token
text = self.tokenizer.decode([int(next_token.item())], skip_special_tokens=True)
if text:
yield text
|