update clip_eval to accept dmx model

clip_eval.py

@@ -29,7 +29,6 @@ _CITATION = """
 }
 """
 
-
 @add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
 class DmxClipEval(evaluate.Metric):
     def _info(self):
@@ -38,19 +37,17 @@ class DmxClipEval(evaluate.Metric):
             description=_DESCRIPTION,
             citation=_CITATION,
             inputs_description=_KWARGS_DESCRIPTION,
-            features=[
-                datasets.Features(
-                    {
-                        "model_name": datasets.Value("string"),
-                        "dataset_names": datasets.Value("string"),
-                        "n_examples": datasets.Value("int32"),
-                    }
-                ),
-            ],
+            features=datasets.Features(
+                {
+                    "model": datasets.Value("string"),
+                    "dataset_names": datasets.Value("string"),
+                    "n_examples": datasets.Value("int32"),
+                }
+            ),
         )
 
     def clip_dataset_evaluator(
-        self, model, device, desc, dataset_name="mscoco", n_examples=-1
+        self, model, device, dataset_name="mscoco", n_examples=-1
     ):
         processor = CLIPProcessor.from_pretrained(model.config._name_or_path)
         if dataset_name == "mscoco":
@@ -116,34 +113,36 @@ class DmxClipEval(evaluate.Metric):
         }
         return metrics
 
-    def clip_evaluator(self, model, device, desc, n_examples=-1):
+    def clip_evaluator(self, model, device, n_examples=-1):
         metrics = {}
-        for name in ["mscoco", "flickr"]:
+        for dataset_name in ["mscoco", "flickr"]:
             metrics.update(
-                self.clip_dataset_evaluator(model, device, desc, name, n_examples)
+                self.clip_dataset_evaluator(model, device, dataset_name, n_examples)
             )
         return metrics
 
-    def _compute(self, model_name, dataset_names, n_examples):
-
-        actual_model_name = model_name[0]
-        actual_dataset_name_str = dataset_names[0]
-        actual_n_examples = n_examples[0]
+    def _compute(self, model, dataset_names, n_examples, **kwargs):
+        dataset = dataset_names[0] 
+        num_examples = n_examples[0]
+        model_input = model[0] 
 
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = CLIPModel.from_pretrained(actual_model_name).to(device)
-
-        datasets_to_evaluate = [actual_dataset_name_str]
-
+        
+        if isinstance(model_input, str):
+            actual_model = CLIPModel.from_pretrained(model_input).to(device)
+        else:
+            actual_model = model_input
+        
+        datasets_to_evaluate = [dataset]
+        
         metrics = {}
-        for ds_name_loop_var in datasets_to_evaluate:
+        for ds_name in datasets_to_evaluate:
             dataset_metrics = self.clip_dataset_evaluator(
-                model=model,
+                model=actual_model,
                 device=device,
-                desc=actual_model_name,
-                dataset_name=ds_name_loop_var,
-                n_examples=actual_n_examples,
+                dataset_name=ds_name,
+                n_examples=num_examples,
             )
             metrics.update(dataset_metrics)
 
-        return metrics
+        return metrics