improved logging

ls1intum · Jan 29, 2025 · d62e3ed · d62e3ed
1 parent c6ec343
commit d62e3ed
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/llm_core/llm_core/models/__init__.py b/llm_core/llm_core/models/__init__.py
@@ -3,7 +3,6 @@
 from langchain.base_language import BaseLanguageModel
 
 from llm_core.models.model_config import ModelConfig
-from athena.logger import logger
 
 DefaultModelConfig: Type[ModelConfig]
 MiniModelConfig: ModelConfig

diff --git a/llm_core/llm_core/models/ollama.py b/llm_core/llm_core/models/ollama.py
@@ -71,7 +71,7 @@ def max_tokens_must_be_positive(cls, v):
                     return v
 
                 def get_model(self) -> BaseLanguageModel:
-                    print("Getting Model: ", self.model_name.value)
+                    logger.info("Getting Model: ", self.model_name.value)
                     """Get the model from the configuration.
 
                     Returns:
@@ -110,7 +110,7 @@ def get_model(self) -> BaseLanguageModel:
                 class Config:
                     title = 'Ollama'
     except Timeout:
-        print("Connection timed out. Skipping server connection step.")
+        logger.info("Connection timed out. Skipping gpu server connection step.")
     except RequestException as e:
-        print(f"Failed to connect to the server: {e}. Skipping this step.")
+        logger.info(f"Failed to connect to the gpu server: {e}. Skipping this step.")