refine eval (#282)

Author: wenhuach21

auto_round/__main__.py

@@ -328,16 +328,22 @@ def tune(args):
         logger.info(f"Using lm-eval version {lm_eval_version}")
         model_args = f"pretrained={eval_folder}"
         model_args = model_args + f",trust_remote_code={not args.disable_trust_remote_code}"
-        user_model = None
         if args.act_bits <= 8:
-            user_model = model.to(device_str)
-
-        res = simple_evaluate(
-            model="hf",
-            model_args=model_args,
-            tasks=tasks,
-            batch_size=args.eval_bs,
-            user_model=user_model)
+            if hasattr(model, "hf_device_map") and len(model.hf_device_map) > 1:
+                from accelerate.big_modeling import dispatch_model
+
+                dispatch_model(model, model.hf_device_map)
+                user_model = model
+            else:
+                user_model = model.to(device_str)
+            if args.eval_bs == "auto":
+                args.eval_bs = 16
+            from auto_round.eval.evaluation import simple_evaluate_user_model
+            res = simple_evaluate_user_model(user_model, tokenizer, tasks=tasks, batch_size=args.eval_bs)
+        else:
+            res = simple_evaluate(model="hf", model_args=model_args,
+                                  tasks=tasks,
+                                  batch_size=args.eval_bs)
         print(make_table(res))
 
 

auto_round/eval/evaluation.py

@@ -23,11 +23,28 @@
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+from lm_eval.models.huggingface import HFLM
+
+
+def simple_evaluate_user_model(
+        user_model,
+        tokenizer,
+        batch_size: Optional[int] = None,
+        max_batch_size: Optional[int] = None,
+        **kwargs
+):
+    hflm = HFLM(pretrained=user_model, tokenizer=tokenizer, batch_size=batch_size, max_batch_size=max_batch_size)
+    return lm_simple_evaluate(
+        model=hflm,
+        model_args=None,
+        batch_size=batch_size,
+        max_batch_size=max_batch_size,
+        **kwargs)
+
 
 def simple_evaluate(
         model,
         model_args: Optional[Union[str, dict]] = None,
-        user_model=None,
         batch_size: Optional[int] = None,
         max_batch_size: Optional[int] = None,
         device: Optional[str] = None,
@@ -37,32 +54,8 @@ def simple_evaluate(
     except:
         from auto_round.auto_quantizer import AutoHfQuantizer
 
-    if model_args is None:
-        model_args = ""
-
-    if isinstance(model_args, dict):
-        lm = lm_eval.api.registry.get_model(model).create_from_arg_obj(
-            model_args,
-            {
-                "batch_size": batch_size,
-                "max_batch_size": max_batch_size,
-                "device": device,
-            },
-        )
-
-    else:
-        lm = lm_eval.api.registry.get_model(model).create_from_arg_string(
-            model_args,
-            {
-                "batch_size": batch_size,
-                "max_batch_size": max_batch_size,
-                "device": device,
-            },
-        )
-    if user_model is not None:
-        lm._model = user_model
     return lm_simple_evaluate(
-        model=lm,
+        model=model,
         model_args=model_args,
         batch_size=batch_size,
         max_batch_size=max_batch_size,

examples/language-modeling/main.py

@@ -446,11 +446,21 @@
         model_args = model_args + f",trust_remote_code={not args.disable_trust_remote_code}"
         user_model = None
         if args.act_bits <= 8:
-            user_model = model.to(device_str)
-
-        res = simple_evaluate(model="hf", model_args=model_args,
-                              tasks=tasks,
-                              batch_size=args.eval_bs, user_model=user_model)
+            if hasattr(model, "hf_device_map") and len(model.hf_device_map) > 1:
+                from accelerate.big_modeling import dispatch_model
+
+                dispatch_model(model, model.hf_device_map)
+                user_model = model
+            else:
+                user_model = model.to(device_str)
+            if args.eval_bs == "auto":
+                args.eval_bs = 16
+            from auto_round.eval.evaluation import  simple_evaluate_user_model
+            res = simple_evaluate_user_model(user_model, tokenizer,tasks=tasks,batch_size=args.eval_bs)
+        else:
+            res = simple_evaluate(model="hf", model_args=model_args,
+                                  tasks=tasks,
+                                  batch_size=args.eval_bs)
         from lm_eval.utils import make_table
 
         print(make_table(res))