File size: 2,923 Bytes
a8639ac |
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 88 89 90 91 92 |
import torch
# from architecture import DecoderTransformer
from builtin_architecture import make_model
import os
import sys
import time
from dataset import dataset, get_train_dataset
import torch.nn.functional as F
EXPERIMENT_DIRECTORY = "runs/code-decoder-v10-vanilla-smaller-batchfirst" # "runs/code-decoder-v9-vanilla-smaller"#"runs/code-decoder-v8-smaller" # "runs/code-decoder-v4-improved" # shakespeare-test, run1-python
device = "mps" if torch.backends.mps.is_available() else "cpu"
device = "cpu"
# net = DecoderTransformer(vocab_size=199, num_blocks=1)
net = make_model()
net.to(device)
net.load_state_dict(
torch.load(os.path.join(EXPERIMENT_DIRECTORY, "ckpt", "best.pt"), weights_only=True)
)
for name, param in net.named_parameters():
if torch.isnan(param).any():
print(f"NaN found in {name}")
for name, param in net.named_parameters():
if param.grad is not None and torch.isnan(param.grad).any():
print(f"NaN found in gradients of {name}")
pad_token_id = 0
sep_token_id = None
input_text = input("Prompt: ")
max_length = 100
input_ids = torch.tensor(dataset.manager.encode(input_text), dtype=int)
print(input_ids.shape)
attention_mask = dataset.manager.attention_mask(input_ids.squeeze(0)).to(device)
generated_text = dataset.manager.decode(input_ids)
print(generated_text)
generated_text = ""
input_ids = torch.randint(199, (1, 1), dtype=torch.long).to(device)
net.eval() # Set model to evaluation mode
temp = 1.0 # Balanced temperature
for _ in range(max_length):
with torch.no_grad():
output = net(input_ids) # Model output
logits = F.log_softmax(output[-1], dim=-1) # Normalize logits
word_weights = logits.div(temp).cpu() # Scale by temperature
# Top-k sampling
top_k = 10 # Adjust based on your vocabulary size
vocab_size = word_weights.size(0)
top_k = min(top_k, vocab_size) # Ensure top_k is valid
top_probs, top_indices = torch.topk(word_weights, k=top_k)
# Handle edge case: only one valid token
if top_probs.size(0) == 1:
word_idx = top_indices[0] # Directly choose the only available token
else:
sampled_idx = torch.multinomial(top_probs, 1).item()
word_idx = top_indices[sampled_idx]
# Decode and append token
print(word_idx)
predicted_token = dataset.manager.decode(word_idx.item())
print(predicted_token, end=" ")
generated_text += predicted_token
print("Word Weights:", word_weights)
print("Top Probabilities:", top_probs)
print("Top Indices:", top_indices)
# Update input sequence
word_tensor = torch.tensor([[word_idx]], dtype=torch.long).to(device)
input_ids = torch.cat([input_ids, word_tensor], dim=1)
print("\nGenerated text:", generated_text)
with open("output.txt", "w+") as f:
f.write(generated_text)
|