improve lm_eval_harness code (#857)

intel · Apr 19, 2023 · c7c557c · c7c557c
1 parent 93c2fb8
commit c7c557c
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 6 deletions.
diff --git a/intel_extension_for_transformers/evaluation/__init__.py b/intel_extension_for_transformers/evaluation/__init__.py
@@ -13,4 +13,5 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License.
+# limitations under the License.
+from .lm_evaluation_harness.evaluator import evaluate
diff --git a/intel_extension_for_transformers/evaluation/lm_evaluation_harness/evaluator.py b/intel_extension_for_transformers/evaluation/lm_evaluation_harness/evaluator.py
@@ -17,6 +17,7 @@
 
 import os
 import random
+import time
 import numpy as np
 import json
 
@@ -134,7 +135,10 @@ def evaluate(model,
 
     dumped = json.dumps(results, indent=2)
     if output_dir:
-        with open(os.path.abspath(os.path.expanduser(output_dir)) + "/results.json", "w") as f:
+        now = int(round(time.time()*1000))
+        date = now02 = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(now/1000))
+        results_name = model_args.split(",")[0].replace("/", "-") + "-" + date + "-"+ "results.json"
+        with open(os.path.abspath(os.path.expanduser(output_dir)) + "/" + results_name, "w") as f:
             f.write(dumped)
 
     return results
diff --git a/intel_extension_for_transformers/evaluation/lm_evaluation_harness/models/huggingface.py b/intel_extension_for_transformers/evaluation/lm_evaluation_harness/models/huggingface.py
@@ -400,9 +400,10 @@ def _create_auto_tokenizer(
     def _model_call(
         self, inputs: TokenSequence, labels: Optional[TokenSequence] = None
     ) -> TokenSequence:
-        if isinstance(self.model, torch.jit._script.RecursiveScriptModule):
-            return self.model(inputs)[0]
-        return self.model(inputs)["logits"]
+        output = self.model(inputs)
+        if isinstance(output, tuple):
+            return output[0]
+        return output["logits"]
 
     def _model_generate(
         self,

diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py
@@ -2,7 +2,7 @@
 import unittest
 import shutil
 import torch
-from intel_extension_for_transformers.evaluation.lm_evaluation_harness.evaluator import evaluate
+from intel_extension_for_transformers.evaluation import evaluate
 from transformers import (
     AutoModelForCausalLM
 )