File size: 2,959 Bytes
37d8122 | 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 | """Génération autoregressive pour Cortex-0 — échantillonnage écrit à la main.
Implémente :
* température (recherche du logits / T) ;
* top-k (garder les k logits les plus grands) ;
* top-p / nucleus (garder le plus petit ensemble de probabilité >= p).
Aucune fonction `model.generate` n'est utilisée : la boucle est explicite.
"""
from __future__ import annotations
import torch
from src.model.text_model import CortexTextLM
from src.tokenizer import CortexTokenizer
from src.tokenizer.special_tokens import BOS_ID, EOS_ID
def _apply_top_k(logits: torch.Tensor, k: int) -> torch.Tensor:
if k is None or k <= 0:
return logits
k = min(k, logits.size(-1))
kth_value = torch.topk(logits, k, dim=-1).values[..., -1, None]
return logits.masked_fill(logits < kth_value, float("-inf"))
def _apply_top_p(logits: torch.Tensor, p: float) -> torch.Tensor:
if p is None or p >= 1.0:
return logits
sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
probs = torch.softmax(sorted_logits, dim=-1)
cumulative = torch.cumsum(probs, dim=-1)
# On retire les tokens au-delà du seuil (en gardant le premier qui dépasse p).
remove = cumulative - probs > p
sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
# Remise dans l'ordre d'origine.
return logits.scatter(-1, sorted_idx, sorted_logits)
@torch.no_grad()
def generate(
model: CortexTextLM,
input_ids: torch.Tensor,
max_new_tokens: int = 50,
temperature: float = 1.0,
top_k: int | None = None,
top_p: float | None = None,
greedy: bool = False,
eos_id: int | None = EOS_ID,
) -> torch.Tensor:
"""Génère ``max_new_tokens`` tokens en continuation de ``input_ids`` (B, T)."""
model.eval()
ctx = model.config.ctx_len
for _ in range(max_new_tokens):
idx_cond = input_ids[:, -ctx:]
logits, _ = model(idx_cond)
logits = logits[:, -1, :] # (B, vocab)
if greedy:
next_id = torch.argmax(logits, dim=-1, keepdim=True)
else:
logits = logits / max(temperature, 1e-6)
logits = _apply_top_k(logits, top_k)
logits = _apply_top_p(logits, top_p)
probs = torch.softmax(logits, dim=-1)
next_id = torch.multinomial(probs, num_samples=1)
input_ids = torch.cat([input_ids, next_id], dim=1)
if eos_id is not None and (next_id == eos_id).all():
break
return input_ids
def generate_text(
model: CortexTextLM,
tokenizer: CortexTokenizer,
prompt: str,
max_new_tokens: int = 50,
**kwargs,
) -> str:
"""Enveloppe pratique : encode le prompt, génère, décode la suite."""
ids = [BOS_ID, *tokenizer.encode(prompt)]
inp = torch.tensor([ids], dtype=torch.long)
out = generate(model, inp, max_new_tokens=max_new_tokens, **kwargs)
return tokenizer.decode(out[0].tolist(), skip_special_tokens=True)
|