Evaluate LLMs with Hugging Face Lighteval on Amazon SageMaker

!pip install sagemaker --upgrade --quiet

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it not exists
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")

import os
import requests as r

lighteval_version = "0.2.0"

# create scripts directory if not exists
os.makedirs("scripts", exist_ok=True)

# load custom scripts from git
raw_github_url = f"https://raw.githubusercontent.com/huggingface/lighteval/v{lighteval_version}/run_evals_accelerate.py"
res = r.get(raw_github_url)
with open("scripts/run_evals_accelerate.py", "w") as f:
    f.write(res.text)

# write requirements.txt
with open("scripts/requirements.txt", "w") as f:
    f.write(f"lighteval=={lighteval_version}")

with open("scripts/tasks.txt", "w") as f:
    f.write(f"lighteval|truthfulqa:mc|0|0")

from sagemaker.huggingface import HuggingFace

# hyperparameters, which are passed into the training job
hyperparameters = {
  'model_args': "pretrained=HuggingFaceH4/zephyr-7b-beta", # Hugging Face Model ID
  'task': 'tasks.txt',  # 'lighteval|truthfulqa:mc|0|0',
  'model_dtype': 'bfloat16', # Torch dtype to load model weights
  'output_dir': '/opt/ml/model' # Directory, which sagemaker uploads to s3 after training
}

# create the Estimator
huggingface_estimator = HuggingFace(
    entry_point          = 'run_evals_accelerate.py',      # train script
    source_dir           = 'scripts',         # directory which includes all the files needed for training
    instance_type        = 'ml.g5.4xlarge',   # instances type used for the training job
    instance_count       = 1,                 # the number of instances used for training
    base_job_name        = "lighteval",          # the name of the training job
    role                 = role,              # Iam role used in training job to access AWS ressources, e.g. S3
    volume_size          = 300,               # the size of the EBS volume in GB
    transformers_version = '4.36',            # the transformers version used in the training job
    pytorch_version      = '2.1',            # the pytorch_version version used in the training job
    py_version           = 'py310',            # the python version used in the training job
    hyperparameters      =  hyperparameters,
    environment          = {
                            "HUGGINGFACE_HUB_CACHE": "/tmp/.cache",
                            # "HF_TOKEN": "REPALCE_WITH_YOUR_TOKEN" # needed for private models
                            }, # set env variable to cache models in /tmp
)

# starting the train job with our uploaded datasets as input
huggingface_estimator.fit()

import tarfile
import json
import io
import os
from sagemaker.s3 import S3Downloader


# download results from s3
results_tar = S3Downloader.read_bytes(huggingface_estimator.model_data)
model_id = hyperparameters["model_args"].split("=")[1]
result={}

# Use tarfile to open the tar content directly from bytes
with tarfile.open(fileobj=io.BytesIO(results_tar), mode="r:gz") as tar:
    # Iterate over items in tar archive to find your json file by its path
    for member in tar.getmembers():
        # get path of results based on model id used to evaluate
        if os.path.join("details", model_id) in member.name and member.name.endswith('.json'):
            # Extract the file content
            f = tar.extractfile(member)
            if f is not None:
                content = f.read()
                result = json.loads(content)
                break

# print results
print(result["results"])
# {'lighteval|truthfulqa:mc|0': {'truthfulqa_mc1': 0.40636474908200737, 'truthfulqa_mc1_stderr': 0.017193835812093897, 'truthfulqa_mc2': 0.5747003398184238, 'truthfulqa_mc2_stderr': 0.015742356478301463}}