import sparseml.transformers

original_model_name = "facebook/opt-125m"
output_directory = "output/"
final_model_name = "nm-testing/opt-125m-pruned2.4"

dataset = "open_platypus"

recipe = """
test_stage:
  obcq_modifiers:
    SparseGPTModifier:
      sparsity: 0.5
      sequential_update: true
      quantize: false
      mask_structure: '2:4'
      targets: ['re:model.decoder.layers.\d*$']
"""

# Apply SparseGPT to the model
sparseml.transformers.oneshot(
    model_name_or_path=original_model_name,
    dataset_name=dataset,
    recipe=recipe,
    output_dir=output_directory,
)

# Upload the output model to Hugging Face Hub
from huggingface_hub import HfApi

HfApi().upload_folder(
    folder_path=output_directory,
    repo_id=final_model_name,
)