File size: 533 Bytes
8496edd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
from .base_agent import BaseAgent
from prompt.template import MODEL_SELECTION_PROMPT
from prompt.constants import modeling_methods
class ModelSelection(BaseAgent):
def __init__(self, llm):
super().__init__(llm)
def select_models(self, modeling_problem: str):
prompt = MODEL_SELECTION_PROMPT.format(modeling_methods=modeling_methods, modeling_problem=modeling_problem)
answer = self.llm.generate(prompt)
selected_models = answer.split('Conclusion:')[-1]
return selected_models
|