Views
No views yet
AutoModel)1from transformers import AutoModel, AutoTokenizer
2
3model = AutoModel.from_pretrained(
4 "flamingo44333/mini-gpt-shakespeare",
5 trust_remote_code=True
6)
7
8tokenizer = AutoTokenizer.from_pretrained(
9 "flamingo44333/mini-gpt-shakespeare",
10 trust_remote_code=True
11)1import torch
2import torch.nn.functional as F
3
4def generate(
5 model,
6 tokenizer,
7 prompt,
8 max_new_tokens=100,
9 temperature=0.5,
10 top_k=40,
11 device="cuda" if torch.cuda.is_available() else "cpu"
12):
13 model.eval()
14 model.to(device)
15
16 # Encode (match training behavior)
17 input_ids = tokenizer.encode(prompt, add_special_tokens=False)
18 input_ids = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0).to(device)
19
20 # Handle DataParallel safely
21 model_to_use = model.module if hasattr(model, "module") else model
22
23 with torch.no_grad():
24 for _ in range(max_new_tokens):
25
26 input_crop = input_ids[:, -model.config.max_seq_len:]
27
28 out = model_to_use(input_crop)
29 logits = out["logits"]
30 pad_id = tokenizer.pad_token_id
31 unk_id = tokenizer.unk_token_id
32 if pad_id is not None:
33 logits[:, -1, pad_id] = float('-inf')
34
35 if unk_id is not None:
36 logits[:, -1, unk_id] = float('-inf')
37 logits = logits[:, -1, :] / temperature
38
39 if top_k is not None:
40 values, indices = torch.topk(logits, top_k)
41 probs = F.softmax(values, dim=-1)
42 next_token = indices.gather(-1, torch.multinomial(probs, 1))
43 else:
44 probs = F.softmax(logits, dim=-1)
45 next_token = torch.multinomial(probs, 1)
46
47 input_ids = torch.cat([input_ids, next_token], dim=1)
48
49 return tokenizer.decode(
50 input_ids[0].tolist(),
51 clean_up_tokenization_spaces=False
52 )trust_remote_code=True to load custom model.generate() API support