Load encoder and hparams, then encode input prompt

Author: satojkovic

pytorch/gpt_pytorch.py

@@ -1,5 +1,9 @@
 import torch
 import argparse
+import sys
+import os
+from utils import load_encoder_hparams
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -17,3 +21,13 @@
 
     state_dict = torch.load(args.model_path)
     print(f"state_dict: {len(state_dict.keys())} params")
+
+    model_size = "124M"
+    models_dir = "models"
+    encoder, hparams = load_encoder_hparams(model_size, models_dir)
+    print("hparams:", hparams)
+
+    print("prompt:", args.prompt)
+    input_ids = encoder.encode(args.prompt)
+    input_text = encoder.decode(input_ids)
+    print("input_ids:", input_ids)

pytorch/utils.py

@@ -0,0 +1,56 @@
+import os
+import json
+import sys
+import requests
+from tqdm import tqdm
+
+# import picoGPT
+sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(os.path.join(os.path.dirname(__file__), "../picoGPT"))
+from encoder import get_encoder
+
+
+# Copy from picoGPT because picoGPT/utils.py import tensorflow
+def download_gpt2_files(model_size, model_dir):
+    assert model_size in ["124M", "355M", "774M", "1558M"]
+    for filename in [
+        "checkpoint",
+        "encoder.json",
+        "hparams.json",
+        "model.ckpt.data-00000-of-00001",
+        "model.ckpt.index",
+        "model.ckpt.meta",
+        "vocab.bpe",
+    ]:
+        url = "https://openaipublic.blob.core.windows.net/gpt-2/models"
+        r = requests.get(f"{url}/{model_size}/{filename}", stream=True)
+        r.raise_for_status()
+
+        with open(os.path.join(model_dir, filename), "wb") as f:
+            file_size = int(r.headers["content-length"])
+            chunk_size = 1000
+            with tqdm(
+                ncols=100,
+                desc="Fetching " + filename,
+                total=file_size,
+                unit_scale=True,
+                unit="b",
+            ) as pbar:
+                # 1k for chunk_size, since Ethernet packet size is around 1500 bytes
+                for chunk in r.iter_content(chunk_size=chunk_size):
+                    f.write(chunk)
+                    pbar.update(chunk_size)
+
+
+def load_encoder_hparams(model_size, models_dir):
+    assert model_size in ["124M", "355M", "774M", "1558M"]
+
+    model_dir = os.path.join(models_dir, model_size)
+    if not model_dir:  # download files if necessary
+        os.makedirs(model_dir, exist_ok=True)
+        download_gpt2_files(model_size, model_dir)
+
+    encoder = get_encoder(model_size, models_dir)
+    hparams = json.load(open(os.path.join(model_dir, "hparams.json")))
+
+    return encoder, hparams