add device

dmxMetric.py

@@ -3,6 +3,7 @@ import lm_eval
 from typing import Union, List, Optional
 from dmx.compressor.dmx import config_rules, DmxModel
 import datasets
+import torch
 
 _DESCRIPTION = """
 Evaluation function using lm-eval with d-Matrix integration.
@@ -54,6 +55,7 @@ class  DmxMetric(evaluate.Metric):
         batch_size: Optional[Union[int, str]] = None,
         max_batch_size: Optional[int] = None,
         limit: Optional[Union[int, float]] = None,
+        device: Optional[str] = None,
         revision: str = "main",
         trust_remote_code: bool = False, 
         log_samples: bool = True,
@@ -63,7 +65,10 @@ class  DmxMetric(evaluate.Metric):
         """
         Evaluate a model on multiple tasks and metrics using lm-eval with optional d-Matrix integration.
         """
-        model_args = f"pretrained={model},revision={revision},trust_remote_code={str(trust_remote_code)}"
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        model_args = f"pretrained={model},revision={revision},trust_remote_code={str(trust_remote_code)},device={device}"
 
         lm = lm_eval.api.registry.get_model("hf").create_from_arg_string(
             model_args,