Adapt integration test to new Pipeline

src/lighteval/models/tgi_model.py

@@ -24,6 +24,7 @@
 from huggingface_hub import AsyncInferenceClient, InferenceClient
 from transformers import AutoTokenizer
 
+from lighteval.models.abstract_model import ModelInfo
 from lighteval.models.endpoint_model import InferenceEndpointModel
 
 
@@ -45,15 +46,19 @@ def __init__(self, address, auth_token=None, model_id=None) -> None:
         self.client = InferenceClient(base_url=address, headers=headers, timeout=240)
         self.async_client = AsyncInferenceClient(base_url=address, headers=headers, timeout=240)
         self._max_gen_toks = 256
-        self.model_info = requests.get(f"{address}/info", headers=headers).json()
-        if "model_id" not in self.model_info:
+        info = requests.get(f"{address}/info", headers=headers).json()
+        if "model_id" not in info:
             raise ValueError("Error occured when fetching info: " + str(self.model_info))
-        if model_id:
-            self.model_info["model_id"] = model_id
-        self._tokenizer = AutoTokenizer.from_pretrained(self.model_info["model_id"])
+        self.name = info["model_id"]
+        self.model_info = ModelInfo(
+            model_name=model_id or self.name,
+            model_sha=info["model_sha"],
+            model_dtype=info["model_dtype"] or "default",
+            model_size=-1,
+        )
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model_info.model_name)
         self._add_special_tokens = True
         self.use_async = True
-        self.name = self.model_info["model_id"]
 
     def set_cache_hook(self, cache_hook):
         self.cache_hook = cache_hook

tests/test_base_model.py

@@ -22,15 +22,17 @@
 
 import os
 from typing import Iterator, TypeAlias
+from unittest.mock import patch
 
 import pytest
 from huggingface_hub import ChatCompletionInputMessage
 from transformers import BatchEncoding
 
-from lighteval.evaluator import EvaluationTracker, evaluate
+from lighteval.logging.evaluation_tracker import EvaluationTracker
 from lighteval.metrics.metrics import Metrics
 from lighteval.models.base_model import BaseModel
 from lighteval.models.model_config import BaseModelConfig, EnvConfig
+from lighteval.pipeline import ParallelismManager, Pipeline, PipelineParameters
 from lighteval.tasks.lighteval_task import LightevalTask, LightevalTaskConfig, create_requests_from_tasks
 from lighteval.tasks.requests import (
     Doc,
@@ -120,28 +122,40 @@ def task(self) -> LightevalTask:
     def test_integration(self, task: LightevalTask, base_model: BaseModel, num_fewshot: int, use_chat_template: bool):
         base_model.use_chat_template = use_chat_template
 
+        env_config = EnvConfig(token=TOKEN, cache_dir=CACHE_PATH)
         evaluation_tracker = EvaluationTracker()
+        pipeline_params = PipelineParameters(
+            launcher_type=ParallelismManager.NONE,
+            env_config=env_config,
+            use_chat_template=use_chat_template,
+            override_batch_size=1,
+        )
+
+        with patch("lighteval.pipeline.Pipeline._init_tasks_and_requests"):
+            pipeline = Pipeline(
+                tasks=f"custom|test|{num_fewshot}|0",
+                pipeline_parameters=pipeline_params,
+                evaluation_tracker=evaluation_tracker,
+                model=base_model,
+            )
         task_dict = {"custom|test": task}
         evaluation_tracker.task_config_logger.log(task_dict)
-        requests_dict, docs = create_requests_from_tasks(
+        fewshot_dict = {"custom|test": [(num_fewshot, False)]}
+        pipeline.task_names_list = ["custom|test"]
+        pipeline.task_dict = task_dict
+        pipeline.fewshot_dict = fewshot_dict
+        requests, docs = create_requests_from_tasks(
             task_dict=task_dict,
-            fewshot_dict={"custom|test": [(num_fewshot, False)]},
-            num_fewshot_seeds=0,
+            fewshot_dict=fewshot_dict,
+            num_fewshot_seeds=pipeline_params.num_fewshot_seeds,
             lm=base_model,
-            max_samples=1,
+            max_samples=pipeline_params.max_samples,
             evaluation_tracker=evaluation_tracker,
             use_chat_template=use_chat_template,
-            system_prompt=None,
+            system_prompt=pipeline_params.system_prompt,
         )
+        pipeline.requests = requests
+        pipeline.docs = docs
+        evaluation_tracker.task_config_logger.log(task_dict)
 
-        evaluation_tracker = evaluate(
-            lm=base_model,
-            requests_dict=requests_dict,
-            docs=docs,
-            task_dict=task_dict,
-            override_bs=1,
-            evaluation_tracker=evaluation_tracker,
-        )
-        evaluation_tracker.metrics_logger.aggregate(task_dict=task_dict)
-        evaluation_tracker.details_logger.aggregate()
-        evaluation_tracker.generate_final_dict()
+        pipeline.evaluate()

tests/test_endpoint_model.py

@@ -25,22 +25,25 @@
 import time
 from collections import defaultdict
 from typing import Iterator, TypeAlias
+from unittest.mock import patch
 
 import docker
 import docker.errors
 import pytest
 import requests
 from huggingface_hub import ChatCompletionInputMessage
 
-from lighteval.evaluator import EvaluationTracker, evaluate
+from lighteval.logging.evaluation_tracker import EvaluationTracker
 from lighteval.metrics.metrics import Metrics
 from lighteval.models.tgi_model import ModelClient as TGIModel
+from lighteval.pipeline import ParallelismManager, Pipeline, PipelineParameters
 from lighteval.tasks.lighteval_task import LightevalTask, LightevalTaskConfig, create_requests_from_tasks
 from lighteval.tasks.requests import (
     Doc,
     Request,
     RequestType,
 )
+from lighteval.utils.utils import EnvConfig
 
 
 TOKEN = os.environ.get("HF_TOKEN")
@@ -164,28 +167,39 @@ def test_loglikelihood(self, zero_shot_request_dict: RequestDict, tgi_model: TGI
     @pytest.mark.parametrize("num_fewshot", [0, 2])
     @pytest.mark.parametrize("use_chat_template", [False, True])
     def test_integration(self, task: LightevalTask, tgi_model: TGIModel, num_fewshot: int, use_chat_template: bool):
+        env_config = EnvConfig(token=TOKEN, cache_dir=CACHE_PATH)
         evaluation_tracker = EvaluationTracker()
+        pipeline_params = PipelineParameters(
+            launcher_type=ParallelismManager.NONE,
+            env_config=env_config,
+            use_chat_template=use_chat_template,
+        )
+
+        with patch("lighteval.pipeline.Pipeline._init_tasks_and_requests"):
+            pipeline = Pipeline(
+                tasks=f"custom|test|{num_fewshot}|0",
+                pipeline_parameters=pipeline_params,
+                evaluation_tracker=evaluation_tracker,
+                model=tgi_model,
+            )
         task_dict = {"custom|test": task}
         evaluation_tracker.task_config_logger.log(task_dict)
-        requests_dict, docs = create_requests_from_tasks(
+        fewshot_dict = {"custom|test": [(num_fewshot, False)]}
+        pipeline.task_names_list = ["custom|test"]
+        pipeline.task_dict = task_dict
+        pipeline.fewshot_dict = fewshot_dict
+        requests, docs = create_requests_from_tasks(
             task_dict=task_dict,
-            fewshot_dict={"custom|test": [(num_fewshot, False)]},
-            num_fewshot_seeds=0,
+            fewshot_dict=fewshot_dict,
+            num_fewshot_seeds=pipeline_params.num_fewshot_seeds,
             lm=tgi_model,
-            max_samples=1,
+            max_samples=pipeline_params.max_samples,
             evaluation_tracker=evaluation_tracker,
             use_chat_template=use_chat_template,
-            system_prompt=None,
+            system_prompt=pipeline_params.system_prompt,
         )
+        pipeline.requests = requests
+        pipeline.docs = docs
+        evaluation_tracker.task_config_logger.log(task_dict)
 
-        evaluation_tracker = evaluate(
-            lm=tgi_model,
-            requests_dict=requests_dict,
-            docs=docs,
-            task_dict=task_dict,
-            override_bs=1,
-            evaluation_tracker=evaluation_tracker,
-        )
-        evaluation_tracker.metrics_logger.aggregate(task_dict=task_dict)
-        evaluation_tracker.details_logger.aggregate()
-        evaluation_tracker.generate_final_dict()
+        pipeline.evaluate()