| import torch |
| from modeling_m31 import load_model_bundle |
|
|
| def generate(model, tokenizer, prompt, max_new_tokens=256, temperature=0.7, top_p=0.95): |
| ids = tokenizer.encode(prompt).ids |
| x = torch.tensor([ids], device=next(model.parameters()).device) |
| eos = tokenizer.token_to_id('<eos>') |
| with torch.no_grad(): |
| for _ in range(max_new_tokens): |
| logits = model(x[:, -1024:])[:, -1, :] / max(temperature, 1e-5) |
| probs = torch.softmax(logits, dim=-1) |
| vals, idx = torch.sort(probs, descending=True) |
| c = torch.cumsum(vals, dim=-1); vals[c > top_p] = 0; vals = vals / vals.sum(dim=-1, keepdim=True) |
| nxt = idx.gather(-1, torch.multinomial(vals, 1)); x = torch.cat([x, nxt], dim=1) |
| if eos is not None and int(nxt.item()) == eos: break |
| return tokenizer.decode(x[0].tolist()) |
|
|
