# Default configuration for the prompt optimisation framework

project:
  name: prompt_optimisation
  version: 0.1.0
  description: A framework for optimising prompts for language models

paths:
  data: &data_path /ADD/PATH/TO/DATA

logging:
  level: INFO
  format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"

experiment:
  name: default_experiment
  description: Default experiment configuration

components:
  llm:
    default: OpenAI
    OpenAI:
      version_name: gpt-4o_v2024-05-13
      temperature: 0.5
      top_p: 0.95

  task:
    default: question_answer
    question_answer:
      data_path: *data_path
      dataset_name: gsm8k
      train_test_flag: test
      # [] to run full dataset
      # [1, 2, 3, 4, 5] to run first 5 examples
      id_lst: []
      prompt_msg_template:
        - role: system
          content: >
            Take a deep breath and work on this problem step-by-step.
        - role: user
          content: "{question}"

  prompt_optimiser:
    default: psao
    psao:
      psao_intro_prompt: ""
      psao_struct_ann: "(importance ann_var)"
      r_seed: 42
      optimise_user_prompt_flag: false
      optuna_db_name: prompt_opt_psao_aqua_db.db
      optuna_study_name: aqua_psao
      optuna_n_trials: 5

steps:
  - name: "init_llm"
    type: "component"
    component_type: "llm"
    component_name: "${components.llm.default}"
  
  - name: "init_task"
    type: "component"
    component_type: "task"
    component_name: "${components.task.default}"
  
  - name: "init_prompt_optimiser"
    type: "component"
    component_type: "prompt_optimiser"
    component_name: "${components.prompt_optimiser.default}"
  
  - name: "run_base_prompt"
    type: "function"
    function: "run_base_prompt"
    depends_on: ["init_llm", "init_task"]
  
  - name: "optimise_prompt"
    type: "function"
    function: "optimise_prompt"
    depends_on: ["init_llm", "init_task", "init_prompt_optimiser"]
  
  - name: "run_generation"
    type: "function"
    function: "run_generation"
    depends_on: ["init_llm", "init_task", "optimise_prompt"]
  
  - name: "evaluate_results"
    type: "function"
    function: "evaluate_results"
    depends_on: ["run_base_prompt", "run_optimised_prompt"]