Add placeholder wrappers for ET evaluation: Eager and Runtime (#2821)

Summary:
Pull Request resolved: #2821

ET Evaluation can reuse most of the work from evaluating AOT pre-lowered models.

This diff introduces 2 skeleton implementations of EvaluationWrappers:
* one for eager
* one for runtime.

Notably, these wrappers are used when a .pte file is provided (with the latter requiring a tokenizer binary)

Reviewed By: larryliu0820

Differential Revision: D55674251

fbshipit-source-id: f972a1728b8c3d5cead2fcf2ad3196f6b3254dd6

Author: Jack-Khuu

examples/models/llama2/eval_llama_lib.py

@@ -84,6 +84,58 @@ def _model_generate(self, context, max_length, eos_token_id):
         raise Exception("unimplemented")
 
 
+class ETEagerEvalWrapper(GPTFastEvalWrapper):
+    """
+    A wrapper class for ExecuTorch Eager integration with the
+    lm-evaluation-harness library.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        tokenizer: SentencePieceProcessor,
+        max_seq_length: Optional[int] = None,
+    ):
+        super().__init__(None, tokenizer, max_seq_length)
+        self._model = model
+
+    def _model_call(self, inps):
+        # Given inps (tokens), return the logits from a single
+        # forward call
+
+        # Example:
+        # inps: Tensor of shape (1, N)
+        # logits: Tensor of shape (1, N, 32000)
+        pass
+
+
+class ETRunnerEvalWrapper(GPTFastEvalWrapper):
+    """
+    A wrapper class for ExecuTorch Runtime integration with the
+    lm-evaluation-harness library.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        tokenizer: SentencePieceProcessor,
+        tokenizer_bin: str,
+        max_seq_length: Optional[int] = None,
+    ):
+        super().__init__(None, tokenizer, max_seq_length)
+        self._model = model
+        self._tokenizer_bin = tokenizer_bin
+
+    def _model_call(self, inps):
+        # Given inps (tokens), return the logits from a single
+        # forward call
+
+        # Example:
+        # inps: Tensor of shape (1, N)
+        # logits: Tensor of shape (1, N, 32000)
+        pass
+
+
 @torch.no_grad()
 def eval(
     eval_wrapper: LM,
@@ -131,6 +183,24 @@ def gen_eval_wrapper(
     """
     tokenizer = SentencePieceProcessor(model_file=str(args.tokenizer_path))
 
+    # ExecuTorch Binary Evaluation
+    if (model := args.pte) is not None:
+        if (tokenizer_bin := args.tokenizer_bin) is not None:
+            # ETRunnerEvalWrapper: Create a wrapper around an ExecuTorch model, evaluated at runtime
+            return ETRunnerEvalWrapper(
+                model=model,
+                tokenizer=tokenizer,
+                tokenizer_bin=tokenizer_bin,
+                max_seq_length=args.max_seq_length,
+            )
+
+        # ETRunnerEvalWrapper: Create a wrapper around an ExecuTorch model, evaluated eagerly
+        return ETEagerEvalWrapper(
+            model=model,
+            tokenizer=tokenizer,
+            max_seq_length=args.max_seq_length,
+        )
+
     # GPTFastEvalWrapper: Create a wrapper around a pre-exported model
     manager: LlamaEdgeManager = _prepare_for_llama_export(model_name, args)
     model = (
@@ -161,6 +231,21 @@ def build_args_parser() -> argparse.ArgumentParser:
         "--limit", type=int, default=5, help="number of samples to evalulate"
     )
 
+    # Add additional args specific to eval via an ET Runner
+    # Note: For initial integration, the tokenizer.model is also required
+    parser.add_argument(
+        "--pte",
+        type=str,
+        default=None,
+        help="[For ExecuTorch] Path to the ExecuTorch model being evaluated. If provided, don't go through the export flow",
+    )
+    parser.add_argument(
+        "--tokenizer_bin",
+        type=str,
+        default=None,
+        help="[For ExecuTorch] Path to the Tokenizer binary for evaluating ExecuTorch models via runtime",
+    )
+
     return parser