File size: 5,208 Bytes
5c7f987 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | import os
import sys
import torch
import sentencepiece as spm
# Add the parent directory of this script to the python path
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from model.config import VaghLMConfig
from model.transformer import VaghLM
def run_generation():
# 1. Initialize configuration and setup device
config = VaghLMConfig()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 2. Load the SentencePiece Tokenizer
tokenizer_path = os.path.join(ROOT, "tokenizer", "vocab", "tokenizer.model")
if not os.path.exists(tokenizer_path):
raise FileNotFoundError(f"Tokenizer not found at {tokenizer_path}. Please check your installation.")
sp = spm.SentencePieceProcessor()
sp.load(tokenizer_path)
print(f"[OK] Tokenizer loaded successfully (Vocab size: {sp.vocab_size()})")
# 3. Initialize Model Architecture
model = VaghLM(config)
# 4. Load the Model Weights (checkpoints)
# Check for finetuned_v2_best.pt first, then fallback to best_checkpoint.pt
checkpoint_names = ["finetuned_v2_best.pt", "best_checkpoint.pt"]
checkpoint_path = None
for ckpt_name in checkpoint_names:
p = os.path.join(ROOT, "checkpoints", ckpt_name)
if os.path.exists(p):
checkpoint_path = p
break
if checkpoint_path is None:
print("[WARNING] No pre-trained checkpoint found in checkpoints/. Initializing with random weights.")
else:
print(f"Loading checkpoint: {checkpoint_path}")
ckpt = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(ckpt["model_state"])
print(f"[OK] Model weights loaded (Step: {ckpt.get('step', 'N/A')}, Loss: {ckpt.get('loss', 'N/A'):.4f})")
model.to(device)
model.eval()
print(f"Model parameters: {model.count_parameters():,}")
# 5. Define Special Tokens
BOS_ID = config.bos_id
USER_ID = config.user_id
ASSISTANT_ID = config.assistant_id
EOS_ID = config.eos_id
# 6. Set Generation Parameters
prompt = "Explain what a black hole is."
max_new_tokens = 150
temperature = 0.7
top_k = 40
top_p = 0.9
repetition_penalty = 1.3
print(f"\nPrompt: {prompt}")
print(f"Generation Parameters: temp={temperature}, top_k={top_k}, top_p={top_p}, rep_penalty={repetition_penalty}")
print("-" * 50)
# 7. Format the input with special tokens
user_tokens = sp.encode(prompt)
input_ids = [BOS_ID, USER_ID] + user_tokens + [ASSISTANT_ID]
x = torch.tensor([input_ids], dtype=torch.long, device=device)
generated = []
# 8. Generation Loop
with torch.no_grad():
for _ in range(max_new_tokens):
# Truncate input if it exceeds the model's max context length
if x.size(1) > config.context_length:
x = x[:, -config.context_length:]
# Forward pass to obtain logits for the next token
logits, _ = model(x)
logits = logits[:, -1, :] # Select only the last sequence position
# Apply Repetition Penalty
for token_id in set(x[0].tolist()):
if logits[0, token_id] > 0:
logits[0, token_id] /= repetition_penalty
else:
logits[0, token_id] *= repetition_penalty
# Apply Temperature
if temperature > 0:
logits = logits / temperature
# Apply Top-k Filtering
if top_k is not None and top_k > 0:
values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < values[:, [-1]]] = -float("inf")
# Apply Top-p (Nucleus) Filtering
if top_p is not None and top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
# Mask tokens that fall outside cumulative probability threshold
sorted_indices_to_remove = (cumulative_probs - torch.softmax(sorted_logits, dim=-1)) > top_p
sorted_logits[sorted_indices_to_remove] = -float("inf")
# Restore original indexing
logits = torch.zeros_like(logits).scatter_(1, sorted_indices, sorted_logits)
# Sample from the probability distribution
if temperature == 0:
next_token = torch.argmax(logits, dim=-1, keepdim=True)
else:
probs = torch.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
x = torch.cat([x, next_token], dim=1)
token_val = next_token.item()
generated.append(token_val)
# Break generation if End Of Sequence token is encountered
if token_val == EOS_ID:
break
# 9. Decode and print output
response = sp.decode(generated)
print(f"\nResponse:\n{response}\n")
if __name__ == "__main__":
run_generation()
|