tweaks and add a simple test

test_all.py

@@ -0,0 +1,53 @@
+"""
+Run simply with
+$ pytest
+"""
+import os
+import pytest # pip install pytest
+import subprocess
+
+import torch
+from model import ModelArgs, Transformer
+
+def test_argmax_inference():
+    """
+    Only the simplest test for now: run inference with temperature 0 
+    (for determinism) in both C and PyTorch, and see that the sampled tokens 
+    are the same.
+    """
+    test_ckpt_dir = "out" # TODO create a dummy test checkpoint for this?
+
+    # run C version
+    model_path = os.path.join(test_ckpt_dir, "model.bin")
+    command = ["./run", model_path, "0.0"]
+    proc = subprocess.Popen(command, stdout=subprocess.PIPE)
+    c_tokens = []
+    for line in proc.stdout:
+        token = int(line.decode('utf-8').strip())
+        c_tokens.append(token)
+    proc.wait()
+    #print(c_tokens)
+
+    # run PyTorch version
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    ckpt_path = os.path.join(test_ckpt_dir, "ckpt.pt")
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    gptconf = ModelArgs(**checkpoint['model_args'])
+    model = Transformer(gptconf)
+    state_dict = checkpoint['model']
+    unwanted_prefix = '_orig_mod.'
+    for k,v in list(state_dict.items()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+    model.to(device)
+    x = torch.tensor([[1]], dtype=torch.long, device=device) # 1 is BOS
+    with torch.inference_mode():
+        y = model.generate(x, max_new_tokens=gptconf.max_seq_len, temperature=0.0)
+    pt_tokens = y[0].tolist()
+    pt_tokens = pt_tokens[1:] # remove BOS
+    #print(pt_tokens)
+
+    # compare
+    assert c_tokens == pt_tokens