Downstream evaluation (#49)

* Initial config for evaluation.

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Basic working implementation of downstream eval using lm_eval_harness.

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Remove unnecessary calls to reset cache

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

---------

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

Author: kaselby

evaluate.py

@@ -0,0 +1,86 @@
+import argparse
+import os
+import logging
+
+import torch
+
+from transformers.trainer_utils import set_seed
+from transformers import AutoConfig, AutoModelForCausalLM
+from lm_eval.models import HFLM
+from lm_eval import simple_evaluate
+from lm_eval.utils import make_table
+from lm_eval.models.huggingface import HFLM
+
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Evaluate trained model on common LM datasets using LM Eval Harness.")
+    parser.add_argument("--model_type", type=str, choices=["hf", "sparse"], default="hf")
+    parser.add_argument("--model_name_or_config", type=str, required=True,
+                       help="Name or path of the base model (e.g., meta-llama/Llama-2-7b-hf)")
+    parser.add_argument("--sp_dir", type=str, default="",
+                        help="Path to trained predictor dir for sparse model.")
+    parser.add_argument("--tasks", nargs='+', default=["hellaswag"], 
+                        help="Tasks on which to evaluate")
+    parser.add_argument("--batch_size", type=int, default=4,
+                       help="Batch size for processing")
+    parser.add_argument("--device", type=str, default="auto",
+                       help="Device to use (auto, cpu, cuda)")
+    return parser
+
+
+def main():
+    args = parse_args()
+
+    # Set seed
+    set_seed(args.seed)
+    
+    # Setup device
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+    
+    logger.info(f"Using device: {device}")
+
+    # Load pretrained model
+    logging.info("Loading pretrained model for evaluation...")
+    config = AutoConfig.from_pretrained(args.model_name_or_config)
+    if args.model_type == "hf":
+        model = AutoModelForCausalLM.from_pretrained(config)
+    if args.model_type == "sparse":
+        model = AutoModelForCausalLM.from_pretrained(config._name_or_path, config=config)
+        for layer_idx, layer in enumerate(model.get_decoder().layers):
+            layer_path = os.path.join(args.sp_dir, f"final_predictor_layer_{layer_idx}")
+            if not os.path.exists(layer_path):
+                logger.error(f"Pretrained weights for sparse predictor at layer {layer_idx} do not exist.")
+                return
+            pretrained_dict = torch.load(layer_path)
+            layer.mlp_lora_proj.load_state_dict(pretrained_dict)
+        model.tie_weights()
+        model.reset_cache()
+
+    wrapped_model = HFLM(
+        pretrained=model,
+        batch_size=args.batch_size,
+        device=device
+    )
+
+    logging.info("Beginning evaluation...")
+    results = simple_evaluate(
+        wrapped_model,
+        tasks=args.tasks,
+        batch_size=args.batch_size,
+        device=device
+    )
+
+    if results is not None:
+        print(make_table(results))
+        if "groups" in results:
+            print(make_table(results, "groups"))
+
+if __name__ == '__main__':
+    main()