Views
No views yet
1import torch
2import torch.nn.functional as F
3# Assuming DiffReaperModel is defined as per train_autogrow.py
4
5def generate(model, tokenizer, prompt, steps=10):
6 model.eval()
7 with torch.no_grad():
8 p_tokens = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
9 p_emb = model.token_embedding(p_tokens[:, :32]) # Hard conditioning
10
11 # Start from pure noise
12 r_noise = torch.randn(1, 32, 1024).to("cuda")
13
14 for i in range(steps):
15 t = torch.tensor([1000 - (i * (1000//steps)) - 1], device="cuda").long()
16 pred = model(torch.cat([p_emb, r_noise], dim=1), t)
17 r_0_pred = pred[:, 32:, :] # Extract response
18 r_noise = 0.4 * r_noise + 0.6 * r_0_pred # Iterative refinement
19
20 # Map to vocab using Cosine Similarity
21 norm_weights = F.normalize(model.token_embedding.weight, dim=-1)
22 norm_r = F.normalize(r_noise, dim=-1)
23 logits = torch.matmul(norm_r, norm_weights.T)
24 return tokenizer.decode(torch.argmax(logits, dim=-1)[0])
25
26# --- Loading Example ---
27# model = DiffReaperModel(vocab_size=50257, n_embd=1024, n_head=16, n_layer=12).to("cuda")
28# model.load_state_dict(torch.load("cropmark_latest.pt"))1 - F.cosine_similarity between predicted and target embeddings.checkpoint_log.txt and uploaded periodically.