valohai.yaml

- step:
    name: data-preprocess
    image: valohai/llm-toolkit:0.2-gpu
    environment: trial2023-g4dn-xlarge
    command:
      - pip install -r requirements-gpu.txt
      - python data-preprocess.py {parameters}
    parameters:
      - name: model_id
        type: string
        default: 'mistralai/Mistral-7B-v0.1'
      - name: max_tokens
        type: integer
        default: 512
    inputs:
      - name: dataset
        default:
        - s3://dd-sample-bucket/mistral/gem-viggo-dataset/test.csv
        - s3://dd-sample-bucket/mistral/gem-viggo-dataset/train.csv
        - s3://dd-sample-bucket/mistral/gem-viggo-dataset/validation.csv

- step:
    name: finetune
    image: valohai/llm-toolkit:0.2-gpu
    environment: trial2023-g4dn-xlarge
    command:
      - pip install -r requirements-gpu.txt
      - python finetune-mistral.py {parameters}
    parameters:
      - name: model_id
        type: string
        default: "mistralai/Mistral-7B-v0.1"
      - name: max_tokens
        type: integer
        default: 512
      - name: output_dir
        type: string
        default: "finetuned_mistral"
      - name: warmup_steps
        type: integer
        default: 5
      - name: max_steps
        type: integer
        default: 30
      - name: learning_rate
        type: float
        default: 2.5e-5
      - name: do_eval
        type: flag
        default: False
    inputs:
      - name: train_data
        default: dataset://viggo/dev_train
      - name: test_data
        default: dataset://viggo/dev_test
      - name: val_data
        default: dataset://viggo/dev_val

- step:
    name: inference
    image: valohai/llm-toolkit:0.2-gpu
    environment: trial2023-g4dn-xlarge
    command:
      - pip install -r requirements-gpu.txt
      - python inference-mistral.py {parameters}
    parameters:
      - name: model_id
        type: string
        default: "mistralai/Mistral-7B-v0.1"
      - name: max_tokens
        type: integer
        default: 512
      - name: sentence
        type: string
        default: Satisfactory is a 2024 game developed by Coffee Stain Studios. It is a great game about building factories on an alien planet.
    inputs:
      - name: finetuned-checkpoint
        default: dataset://mistral-models/best_mistral_checkpoint

- pipeline:
    name: training-pipeline
    nodes:
      - name: preprocess
        type: execution
        step: data-preprocess
      - name: train
        type: execution
        step: finetune
      - name: inference
        type: execution
        step: inference
    edges:
      - [preprocess.output.encoded_val/*, train.input.val_data]
      - [preprocess.output.encoded_train/*, train.input.train_data]
      - [preprocess.output.encoded_test/*, train.input.test_data]
      - [train.output.finetuned_mistral.best_model/*, inference.input.finetuned-checkpoint]
      - [preprocess.parameter.max_tokens, train.parameter.max_tokens]
      - [train.parameter.max_tokens, inference.parameter.max_tokens]
      - [preprocess.parameter.model_id, train.parameter.model_id]
      - [train.parameter.model_id, inference.parameter.model_id]