From d1844938c6c09f3596cc1a9ebcc1d3f5f633f7f7 Mon Sep 17 00:00:00 2001 From: Kye Date: Sat, 29 Jul 2023 11:27:32 -0400 Subject: [PATCH] clean up llm class inputs Former-commit-id: 0269cf13c3fcd9d0dd3773ddcc4cc53f2a61e4ff --- swarms/swarms.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/swarms/swarms.py b/swarms/swarms.py index 15da612a..753b8f40 100644 --- a/swarms/swarms.py +++ b/swarms/swarms.py @@ -70,7 +70,7 @@ class HierarchicalSwarm: - def initialize_llm(self, llm_class): + def initialize_llm(self, llm_class, temperature=0.5): """ Init LLM @@ -80,10 +80,14 @@ class HierarchicalSwarm: """ try: # Initialize language model - self.llm = llm_class(openai_api_key=self.openai_api_key, temperature=self.temperature) + if self.llm_class == OpenAI: + return llm_class(openai_api_key=self.openai_api_key, temperature=self.temperature) + else: + return self.llm_class(model_id="gpt-2", temperature=self.temperature) except Exception as e: logging.error(f"Failed to initialize language model: {e}") + def initialize_tools(self, llm_class, extra_tools=None): """ Init tools