Generate tokens

Author: satojkovic

pytorch/gpt_pytorch.py

@@ -5,6 +5,7 @@
 from utils import load_encoder_hparams_and_params
 from model import GPT2
 import torchsummaryX
+from tqdm import tqdm
 
 
 if __name__ == "__main__":
@@ -31,3 +32,13 @@
 
     model = GPT2(params, hparams, drop_p=0.1)
     torchsummaryX.summary(model, torch.ones(1, len(input_ids), dtype=torch.long))
+
+    for _ in tqdm(range(args.n_tokens_to_generate), "generating"):
+        logits = model(torch.tensor(input_ids).unsqueeze(0))
+        next_id = torch.argmax(logits[0][-1], dim=-1)
+        input_ids.append(next_id.item())
+    print("Input text:\n", input_text)
+    print(
+        "Generated:\n",
+        encoder.decode(input_ids[len(input_ids) - args.n_tokens_to_generate :]),
+    )