!pip install transformers "datasets[s3]==2.18.0" "sagemaker>=2.190.0" "huggingface_hub[cli]" --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}")

{"question": "<question>", "context": "<relevant context to answer>"}
{"question": "<question>", "context": "<relevant context to answer>"}
{"question": "<question>", "context": "<relevant context to answer>"}

from datasets import load_dataset

# Load dataset from the hub
dataset = load_dataset("philschmid/finanical-rag-embedding-dataset", split="train")
input_path = f"s3://{sess.default_bucket()}/datasets/rag-embedding"

# rename columns
dataset = dataset.rename_column("question", "anchor")
dataset = dataset.rename_column("context", "positive")

# Add an id column to the dataset
dataset = dataset.add_column("id", range(len(dataset)))

# split dataset into a 10% test set
dataset = dataset.train_test_split(test_size=0.1)

# save train_dataset to s3 using our SageMaker session

# save datasets to s3
dataset["train"].to_json(f"{input_path}/train/dataset.json", orient="records")
train_dataset_s3_path = f"{input_path}/train/dataset.json"
dataset["test"].to_json(f"{input_path}/test/dataset.json", orient="records")
test_dataset_s3_path = f"{input_path}/test/dataset.json"

print(f"Training data uploaded to:")
print(train_dataset_s3_path)
print(test_dataset_s3_path)
print(
    f"https://s3.console.aws.amazon.com/s3/buckets/{sess.default_bucket()}/?region={sess.boto_region_name}&prefix={input_path.split('/', 3)[-1]}/"
)

from sagemaker.huggingface import HuggingFace

# define Training Job Name
job_name = f"bge-base-exp1"

# define hyperparameters, which are passed into the training job
training_arguments = {
    "model_id": "BAAI/bge-base-en-v1.5",  # model id from the hub
    "train_dataset_path": "/opt/ml/input/data/train/",  # path inside the container where the training data is stored
    "test_dataset_path": "/opt/ml/input/data/test/",  # path inside the container where the test data is stored
    "num_train_epochs": 3,  # number of training epochs
    "learning_rate": 2e-5,  # learning rate
}

# create the Estimator
huggingface_estimator = HuggingFace(
    entry_point="run_mnr.py",  # train script
    source_dir="scripts",  # directory which includes all the files needed for training
    instance_type="ml.g5.xlarge",  # instances type used for the training job
    instance_count=1,  # the number of instances used for training
    max_run=2 * 24 * 60 * 60,  # maximum runtime in seconds (days * hours * minutes * seconds)
    base_job_name=job_name,  # the name of the training job
    role=role,  # Iam role used in training job to access AWS ressources, e.g. S3
    transformers_version="4.36.0",  # the transformers version used in the training job
    pytorch_version="2.1.0",  # the pytorch_version version used in the training job
    py_version="py310",  # the python version used in the training job
    hyperparameters=training_arguments,
    disable_output_compression=True,  # not compress output to save training time and cost
    environment={
        "HUGGINGFACE_HUB_CACHE": "/tmp/.cache",  # set env variable to cache models in /tmp
    },
)

# define a data input dictonary with our uploaded s3 uris
data = {
    "train": train_dataset_s3_path,
    "test": test_dataset_s3_path,
}

# starting the train job with our uploaded datasets as input
huggingface_estimator.fit(data, wait=True)

from sagemaker.huggingface import get_huggingface_llm_image_uri


# retrieve the image uri based on instance type
def get_image_uri(instance_type):
    key = (
        "huggingface-tei"
        if instance_type.startswith("ml.g") or instance_type.startswith("ml.p")
        else "huggingface-tei-cpu"
    )
    return get_huggingface_llm_image_uri(key, version="1.4.0")

from sagemaker.huggingface import HuggingFaceModel

# sagemaker config
instance_type = "ml.c6i.2xlarge"

# create HuggingFaceModel with the image uri
emb_model = HuggingFaceModel(
    role=role,
    image_uri=get_image_uri(instance_type),
    model_data=huggingface_estimator.model_data,
    env={"HF_MODEL_ID": "/opt/ml/model"},  # Path to the model in the container
)

# Deploy model to an endpoint
emb = emb_model.deploy(
    initial_instance_count=1,
    instance_type=instance_type,
)

data = {
    "inputs": "the mesmerizing performances of the leads keep the film grounded and keep the audience riveted .",
}

res = emb.predict(data=data)


# print some results
print(f"length of embeddings: {len(res[0])}")
print(f"first 10 elements of embeddings: {res[0][:10]}")

data = {
    "inputs": "the mesmerizing performances of the leads keep the film grounded and keep the audience riveted .",
}

res = emb.predict(data=data)

# truncate embeddings to matryoshka dimensions
dim = 256
res = res[0][0:dim]

# print some results
print(f"length of embeddings: {len(res)}")

emb.delete_model()
emb.delete_endpoint()