diff --git a/automation-api/yival_experiments/custom_configuration/model_compare.py b/automation-api/yival_experiments/custom_configuration/model_compare.py
index f72875e..fab947b 100644
--- a/automation-api/yival_experiments/custom_configuration/model_compare.py
+++ b/automation-api/yival_experiments/custom_configuration/model_compare.py
@@ -1,11 +1,6 @@
 import litellm
-from alibaba_complete import llm_complete as alibaba_llm_complete
 from litellm import completion
 from model_config_wrapper import ModelConfigWrapper
-from palm_completion import (
-    safety_settings_new_categories,
-    safety_settings_old_categories,
-)
 from yival.logger.token_logger import TokenLogger
 from yival.schemas.experiment_config import MultimodalOutput
 from yival.schemas.model_configs import Response
@@ -14,6 +9,13 @@
 
 # load env vars
 from lib.config import read_config
+from yival_experiments.custom_configuration.llms.alibaba_complete import (
+    llm_complete as alibaba_llm_complete,
+)
+from yival_experiments.custom_configuration.llms.palm_completion import (
+    safety_settings_new_categories,
+    safety_settings_old_categories,
+)
 
 read_config()
 # default model config if not provided