clean up llm class inputs

Former-commit-id: 623eca6eb1
pull/47/head
Kye 2 years ago
parent 620260aa62
commit 361fca9856

@ -70,7 +70,7 @@ class HierarchicalSwarm:
def initialize_llm(self, llm_class): def initialize_llm(self, llm_class, temperature=0.5):
""" """
Init LLM Init LLM
@ -80,10 +80,14 @@ class HierarchicalSwarm:
""" """
try: try:
# Initialize language model # 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: except Exception as e:
logging.error(f"Failed to initialize language model: {e}") logging.error(f"Failed to initialize language model: {e}")
def initialize_tools(self, llm_class, extra_tools=None): def initialize_tools(self, llm_class, extra_tools=None):
""" """
Init tools Init tools

Loading…
Cancel
Save