Spaces:
Build error
Build error
File size: 2,629 Bytes
d660b02 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from pathlib import Path
from huggingface_hub import HfApi
from loguru import logger
try:
from sagemaker.huggingface import HuggingFace
except ModuleNotFoundError:
logger.warning("Couldn't load SageMaker imports. Run 'poetry install --with aws' to support AWS.")
from llm_engineering.settings import settings
finetuning_dir = Path(__file__).resolve().parent
finetuning_requirements_path = finetuning_dir / "requirements.txt"
def run_finetuning_on_sagemaker(
finetuning_type: str = "sft",
num_train_epochs: int = 3,
per_device_train_batch_size: int = 2,
learning_rate: float = 3e-4,
dataset_huggingface_workspace: str = "mlabonne",
is_dummy: bool = False,
) -> None:
assert settings.HUGGINGFACE_ACCESS_TOKEN, "Hugging Face access token is required."
assert settings.AWS_ARN_ROLE, "AWS ARN role is required."
if not finetuning_dir.exists():
raise FileNotFoundError(f"The directory {finetuning_dir} does not exist.")
if not finetuning_requirements_path.exists():
raise FileNotFoundError(f"The file {finetuning_requirements_path} does not exist.")
api = HfApi()
user_info = api.whoami(token=settings.HUGGINGFACE_ACCESS_TOKEN)
huggingface_user = user_info["name"]
logger.info(f"Current Hugging Face user: {huggingface_user}")
hyperparameters = {
"finetuning_type": finetuning_type,
"num_train_epochs": num_train_epochs,
"per_device_train_batch_size": per_device_train_batch_size,
"learning_rate": learning_rate,
"dataset_huggingface_workspace": dataset_huggingface_workspace,
"model_output_huggingface_workspace": huggingface_user,
}
if is_dummy:
hyperparameters["is_dummy"] = True
# Create the HuggingFace SageMaker estimator
huggingface_estimator = HuggingFace(
entry_point="finetune.py",
source_dir=str(finetuning_dir),
instance_type="ml.g5.2xlarge",
instance_count=1,
role=settings.AWS_ARN_ROLE,
transformers_version="4.36",
pytorch_version="2.1",
py_version="py310",
hyperparameters=hyperparameters,
requirements_file=finetuning_requirements_path,
environment={
"HUGGING_FACE_HUB_TOKEN": settings.HUGGINGFACE_ACCESS_TOKEN,
"COMET_API_KEY": settings.COMET_API_KEY,
"COMET_PROJECT_NAME": settings.COMET_PROJECT,
},
)
# Start the training job on SageMaker.
huggingface_estimator.fit()
if __name__ == "__main__":
run_finetuning_on_sagemaker()
|