Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """Setun-Cortex — 采样工具""" | |
| import torch | |
| import torch.nn.functional as F | |
| def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')): | |
| """Filter logits using top-k and/or top-p (nucleus) sampling.""" | |
| if top_k > 0: | |
| top_k = min(top_k, logits.size(-1)) | |
| indices_to_remove = logits < torch.topk(logits, top_k)[0][:, -1:] | |
| logits[indices_to_remove] = filter_value | |
| if top_p > 0.0: | |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) | |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) | |
| sorted_indices_to_remove = cumulative_probs > top_p | |
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() | |
| sorted_indices_to_remove[..., 0] = 0 | |
| indices_to_remove = sorted_indices[sorted_indices_to_remove] | |
| logits[..., indices_to_remove] = filter_value | |
| return logits | |
| def sample_next_token(logits, temperature=0.8, top_k=40, top_p=0.9, repetition_penalty=1.0, prev_tokens=None): | |
| """从 logits 采样下一个 token""" | |
| logits = logits[:, -1, :] / max(temperature, 1e-5) | |
| # Repetition penalty | |
| if repetition_penalty != 1.0 and prev_tokens is not None: | |
| for token in prev_tokens[0].tolist(): | |
| if logits[0, token] < 0: | |
| logits[0, token] *= repetition_penalty | |
| else: | |
| logits[0, token] /= repetition_penalty | |
| # Top-K / Top-P filtering | |
| logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) | |
| probs = F.softmax(logits, dim=-1) | |
| return torch.multinomial(probs, 1) | |
| def generate(model, prompt_ids, max_new_tokens=300, temperature=0.8, | |
| top_k=40, top_p=0.9, repetition_penalty=1.0, seq_len=64, alpha_scale=0.5, | |
| beam_width=1): | |
| """自回归生成文本; beam_width=1 → 采样模式, >1 → beam search (2026-08-05)""" | |
| model.eval() | |
| if beam_width <= 1: | |
| gen = prompt_ids.clone() | |
| for _ in range(max_new_tokens): | |
| ctx = gen[:, -seq_len:] | |
| logits = model(ctx, alpha_scale=alpha_scale) | |
| next_id = sample_next_token(logits, temperature, top_k, top_p, | |
| repetition_penalty, gen[:, -seq_len:] if repetition_penalty > 1 else None) | |
| gen = torch.cat([gen, next_id], dim=1) | |
| return gen | |
| # ===== beam search (v17c 2026-08-05: 消除采样重复死循环) ===== | |
| import math | |
| beams = [(prompt_ids.clone(), 0.0)] | |
| for _ in range(max_new_tokens): | |
| candidates = [] | |
| for seq, score in beams: | |
| ctx = seq[:, -seq_len:] | |
| logits = model(ctx, alpha_scale=alpha_scale)[:, -1, :] / max(temperature, 1e-5) | |
| if repetition_penalty != 1.0: | |
| for t in seq[0].tolist(): | |
| if logits[0, t] < 0: | |
| logits[0, t] *= repetition_penalty | |
| else: | |
| logits[0, t] /= repetition_penalty | |
| logp = F.log_softmax(logits, dim=-1)[0] | |
| topk = torch.topk(logp, beam_width * 2) | |
| for v, lp in zip(topk.indices.tolist(), topk.values.tolist()): | |
| new_seq = torch.cat([seq, torch.tensor([[v]], dtype=torch.long, device=seq.device)], dim=1) | |
| candidates.append((new_seq, score + lp)) | |
| candidates.sort(key=lambda x: x[1], reverse=True) | |
| beams = candidates[:beam_width] | |
| best = max(beams, key=lambda x: x[1] / ((len(x[0][0]) + 5) ** 0.8)) | |
| return best[0] | |