| """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 |
|
|
