From 932fd2c7e6377ffd67a3bd6c451fe4dcc1261a7a Mon Sep 17 00:00:00 2001
From: Joel Niklaus <joel@harvey.ai>
Date: Thu, 12 Dec 2024 17:04:54 +0100
Subject: [PATCH] Fixed import issues.

---
 src/lighteval/main_endpoint.py        | 3 ++-
 src/lighteval/models/litellm_model.py | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/lighteval/main_endpoint.py b/src/lighteval/main_endpoint.py
index 2fa9a61c0..9423d839e 100644
--- a/src/lighteval/main_endpoint.py
+++ b/src/lighteval/main_endpoint.py
@@ -26,6 +26,8 @@
 from typer import Argument, Option
 from typing_extensions import Annotated
 
+from lighteval.models.litellm_model import LiteLLMModelConfig
+
 
 app = typer.Typer()
 
@@ -432,7 +434,6 @@ def litellm(
     """
 
     from lighteval.logging.evaluation_tracker import EvaluationTracker
-    from lighteval.models.model_config import LiteLLMModelConfig
     from lighteval.pipeline import EnvConfig, ParallelismManager, Pipeline, PipelineParameters
 
     env_config = EnvConfig(token=TOKEN, cache_dir=cache_dir)
diff --git a/src/lighteval/models/litellm_model.py b/src/lighteval/models/litellm_model.py
index 22d22d142..d8bd5e86e 100644
--- a/src/lighteval/models/litellm_model.py
+++ b/src/lighteval/models/litellm_model.py
@@ -32,7 +32,7 @@
 
 from lighteval.data import GenerativeTaskDataset
 from lighteval.models.abstract_model import LightevalModel
-from lighteval.models.endpoint_model import ModelInfo
+from lighteval.models.endpoints.endpoint_model import ModelInfo
 from lighteval.models.model_output import (
     GenerativeResponse,
     LoglikelihoodResponse,