Add files using large-upload tool

evaluating_checkpoint.py

@@ -0,0 +1,115 @@
+"""Script to evaluate a model on a validation set. Based on scripts/generate.py from open_lm repo.
+"""
+import argparse
+import json
+import re
+
+import torch
+
+from open_lm.evaluate import evaluate_loop
+from open_lm.data import get_data
+from open_lm.model import create_model
+from open_lm.distributed import init_distributed_device
+from open_lm.params import parse_args
+
+from scripts.generate_without_hf import Generator, GenerationArgs
+
+def generate_model_jsonl(params):
+    params_to_width_depth_dict = {5: (96, 3),
+                    7: (128, 4),
+                    9: (160, 5),
+                    15: (224, 6),
+                    22: (288, 8),
+                    28: (320, 9),
+                    37: (384, 10),
+                    57: (480, 12),
+                    84: (576, 14),
+                    108: (640, 15),
+                    149: (704, 18),
+                    220: (832, 21),
+                    347: (1024, 23),
+                    455: (1120, 26),
+                    611: (1312, 26),
+                    901: (1504, 30)
+                    }
+
+    width, depth = params_to_width_depth_dict[params]
+    filepath = f"layers={depth}_hidden-dim={width}.json"
+    data = {
+        "hidden_dim": width,
+        "n_layers": depth,
+        "n_heads": 4,
+        "seq_len": 2048,
+        "vocab_size": 50432,
+        "post_embed_norm": False,
+        "weight_tying": False,
+        "qk_norm": True
+    }
+
+    with open(filepath, 'w') as file:
+        file.write(json.dumps(data) + '\n')
+    return filepath
+
+
+class ModelArgs:
+    def __init__(self, params, val_data, val_data_key):
+        default_params = vars(parse_args(""))
+        for k, v in default_params.items():
+            setattr(self, k, v)
+        self.model = generate_model_jsonl(params)
+        self.val_data = [val_data]
+        self.val_data_key = [val_data_key]
+        self.per_gpu_val_batch_size = 16
+        self.vocab_size = 50432
+        self.seq_len = 2048
+        self.wandb = False
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", default="path/to/checkpoint")
+
+    parser.add_argument("--val-data", default="", help="Path to validation data. If empty, generate text.")
+    parser.add_argument("--val-data-key", default="json.gz")
+
+    parser.add_argument("--input-text", default="", type=str, help="Input text to generate from. If empty, evaluate on validation data.")
+    parser.add_argument("--max-gen-len", default=200, type=int)
+    parser.add_argument("--temperature", default=0.8, type=float)
+    parser.add_argument("--top-p", default=0.95, type=float)
+
+    args = parser.parse_args()
+    params = int(re.search(r"params=(\d+)", args.checkpoint).group(1))
+
+    checkpoint = torch.load(args.checkpoint)
+    state_dict = checkpoint["state_dict"]
+    state_dict = {x.replace("module.", ""): y for x, y in state_dict.items()}
+    model_args = ModelArgs(params=params, val_data=args.val_data, val_data_key=args.val_data_key)
+    device = init_distributed_device(model_args)
+    model_args.device = device
+    model = create_model(model_args)
+    model.load_state_dict(state_dict)
+    model.eval().cuda()
+    if args.val_data != "":
+        data = get_data(
+            model_args,
+            skip_train=True,
+        )
+        metrics = evaluate_loop(model, data["val_list"], 0, model_args, None)
+        print(metrics)
+    elif args.input_text != "":
+        model = model.half()
+        generator = Generator(model)
+        input_text = [
+            args.input_text,
+        ]
+        output = generator.generate(
+            input_text,
+            GenerationArgs(args.max_gen_len, args.temperature, args.top_p),
+        )
+        print("".join(output))
+        
+    else:
+        print("Please provide either --val-data or --input-text")
+
+if __name__ == "__main__":
+    main()