| import os |
| import torch |
| import torch.nn.functional as F |
| import tiktoken |
| from model import Crimson, MAX_SEQ_LEN |
|
|
| MODEL_PATH = "crimson_8.5M.pt" |
| VOCAB_PATH = "vocab_map.pt" |
| TOKENIZER_NAME = "gpt2" |
|
|
| PAD_ID = 0 |
| SEP_ID = 1 |
| EOS_ID = 2 |
| OFFSET = 3 |
|
|
| def load_model_and_vocab(device): |
| if not os.path.exists(VOCAB_PATH): |
| return None, None, None |
| vocab_data = torch.load(VOCAB_PATH, map_location="cpu") |
| used_tokens = vocab_data["used_tokens"] |
| id2new = vocab_data["id2new"] |
| vocab_size = len(used_tokens) + OFFSET |
| model = Crimson(vocab_size).to(device) |
| if os.path.exists(MODEL_PATH): |
| model.load_state_dict(torch.load(MODEL_PATH, map_location=device)) |
| model.eval() |
| else: |
| return None, None, None |
| return model, used_tokens, id2new |
|
|
| @torch.no_grad() |
| def generate(model, prompt, tokenizer, id2new, used_tokens, device, max_new_tokens=200, temperature=0.8, top_k=50): |
| model.eval() |
| raw_ids = tokenizer.encode(prompt) |
| input_ids = [id2new[rid] for rid in raw_ids if rid in id2new] |
| if not input_ids: input_ids = [PAD_ID] |
| x = torch.tensor([input_ids], dtype=torch.long, device=device) |
| generated = [] |
| for _ in range(max_new_tokens): |
| ctx = x[:, -MAX_SEQ_LEN:] if x.size(1) > MAX_SEQ_LEN else x |
| logits = model(ctx) |
| next_token_logits = logits[:, -1, :] / temperature |
| if top_k is not None: |
| v, _ = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1))) |
| next_token_logits[next_token_logits < v[:, [-1]]] = -float('Inf') |
| probs = F.softmax(next_token_logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
| idx = next_token.item() |
| if idx == EOS_ID: break |
| x = torch.cat((x, next_token), dim=1) |
| generated.append(idx) |
| return tokenizer.decode([used_tokens[i - OFFSET] for i in generated if i >= OFFSET]) |
|
|
| if __name__ == "__main__": |
| device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" |
| model, used_tokens, id2new = load_model_and_vocab(device) |
| enc = tiktoken.get_encoding(TOKENIZER_NAME) |
| if model: |
| newline_id = id2new.get(enc.encode("\n")[0], OFFSET) |
| while True: |
| x = torch.tensor([[newline_id]], dtype=torch.long, device=device) |
| with torch.no_grad(): |
| for _ in range(900): |
| ctx = x[:, -MAX_SEQ_LEN:] if x.size(1) > MAX_SEQ_LEN else x |
| logits = model(ctx) |
| logits = logits[:, -1, :] / 0.8 |
| v, _ = torch.topk(logits, min(50, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float('Inf') |
| probs = F.softmax(logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
| idx = next_token.item() |
| x = torch.cat((x, next_token), dim=1) |
| if idx == EOS_ID: break |
| if idx >= OFFSET: |
| print(enc.decode([used_tokens[idx - OFFSET]]), end="", flush=True) |
| if input("\nPress [Enter] to generate again, or type 'exit': ").lower() == 'exit': break |
|
|
