Enable One-Shot Launch from Finetuning Script

[Satrat]

This PR sets up the finetuning script to be able to launch a one-shot flow. The change log for alternating finetuning was starting to pile up, so wanted to open this as a preliminary PR before adding in the alternating logic.

Summary of Changes

• Add oneshot path to text_generation.py (originally the finetuning script)
• Added StageRunner class for managing the different flows (train, eval, oneshot), rather than doing everything inside the entrypoint text_generation.py script
• Moved a bunch of helper functions that are now shared between text_generation.py and obcq.py (Eventually we can probably get rid of the OBCQ script)

Testing

test_oneshot_recipe.yaml

test_stage:
  obcq_modifiers:
    SparseGPTModifier:
      sparsity: 0.5
      block_size: 128
      sequential_update: False
      quantize: False
      percdamp: 0.01
      prunen: 0
      prunem: 0
      targets: [
        "re:model.layers.\\d+$"
      ]
      target_ids: ["attention_mask", "position_ids"]  

Test script:

def run():
    from sparseml.transformers.finetune.text_generation import run_general
    
    model = "Xenova/llama2.c-stories15M"
    dataset_name = "open_platypus"
    concatenate_data = False
    do_oneshot = True
    output_dir = "./output_oneshot"
    overwrite_output_dir = True
    recipe = "test_oneshot_recipe.yaml"
    splits = {
        "calibration": "train[:90%]",
    }

    run_general(
        model_name_or_path=model,
        dataset_name=dataset_name,
        do_oneshot=do_oneshot,
        output_dir=output_dir,
        overwrite_output_dir=overwrite_output_dir,
        recipe=recipe,
        concatenate_data = concatenate_data,
        splits = splits
    )

if __name__ == "__main__":
    run()