Import a query answering fine-tuned mannequin into Amazon Bedrock as a customized mannequin

Amazon Bedrock is a completely managed service that gives a alternative of high-performing basis fashions (FMs) from main AI corporations like AI21 Labs, Anthropic, Cohere, Meta, Mistral AI, Stability AI, and Amazon via a single API, together with a broad set of capabilities to construct generative AI purposes with safety, privateness, and accountable AI.

Widespread generative AI use instances, together with however not restricted to chatbots, digital assistants, conversational search, and agent assistants, use FMs to offer responses. Retrieval Increase Technology (RAG) is a method to optimize the output of FMs by offering context across the questions for these use instances. Tremendous-tuning the FM is beneficial to additional optimize the output to comply with the model and business voice or vocabulary.

Customized Mannequin Import for Amazon Bedrock, in preview now, lets you import personalized FMs created in different environments, reminiscent of Amazon SageMaker, Amazon Elastic Compute Cloud (Amazon EC2) cases, and on premises, into Amazon Bedrock. This publish is a part of a collection that demonstrates numerous structure patterns for importing fine-tuned FMs into Amazon Bedrock.

On this publish, we offer a step-by-step method of fine-tuning a Mistral mannequin utilizing SageMaker and import it into Amazon Bedrock utilizing the Customized Import Mannequin characteristic. We use the OpenOrca dataset to fine-tune the Mistral mannequin and use the SageMaker FMEval library to judge the fine-tuned mannequin imported into Amazon Bedrock.

Key Options

Among the key options of Customized Mannequin Import for Amazon Bedrock are:

1. This characteristic lets you carry your fine-tuned fashions and leverage the totally managed serverless capabilities of Amazon Bedrock
2. At the moment we’re supporting Llama 2, Llama 3, Flan, Mistral Mannequin architectures utilizing this characteristic with a precisions of FP32, FP16 and BF16 with additional quantizations coming quickly.
3. To leverage this characteristic you possibly can run the import course of (lined later within the weblog) together with your mannequin weights being in Amazon Easy Storage Service (Amazon S3).
4. You’ll be able to even leverage your fashions created utilizing Amazon SageMaker by referencing the Amazon SageMaker mannequin Amazon Useful resource Names (ARN) which gives for a seamless integration with SageMaker.
5. Amazon Bedrock will mechanically scale your mannequin as your site visitors sample will increase and when not in use, scale your mannequin all the way down to 0 thus lowering your prices.

Allow us to dive right into a use-case and see how simple it’s to make use of this characteristic.

Resolution overview

On the time of writing, the Customized Mannequin Import characteristic in Amazon Bedrock helps fashions following the architectures and patterns within the following determine.

On this publish, we stroll via the next high-level steps:

1. Tremendous-tune the mannequin utilizing SageMaker.
2. Import the fine-tuned mannequin into Amazon Bedrock.
3. Take a look at the imported mannequin.
4. Consider the imported mannequin utilizing the FMEval library.

The next diagram illustrates the answer structure.

The method contains the next steps:

1. We use a SageMaker coaching job to fine-tune the mannequin utilizing a SageMaker JupyterLab pocket book. This coaching job reads the dataset from Amazon Easy Storage Service (Amazon S3) and writes the mannequin again into Amazon S3. This mannequin will then be imported into Amazon Bedrock.
2. To import the fine-tuned mannequin, you should use the Amazon Bedrock console, the Boto3 library, or APIs.
3. An import job orchestrates the method to import the mannequin and make the mannequin out there from the client account.
   1. The import job copies all of the mannequin artifacts from the person’s account into an AWS managed S3 bucket.
4. When the import job is full, the fine-tuned mannequin is made out there for invocation out of your AWS account.
5. We use the SageMaker FMEval library in a SageMaker pocket book to judge the imported mannequin.

The copied mannequin artifacts will stay within the Amazon Bedrock account till the customized imported mannequin is deleted from Amazon Bedrock. Deleting mannequin artifacts in your AWS account S3 bucket doesn’t delete the mannequin or the associated artifacts within the Amazon Bedrock managed account. You’ll be able to delete an imported mannequin from Amazon Bedrock together with all of the copied artifacts utilizing both the Amazon Bedrock console, Boto3 library, or APIs.

Moreover, all knowledge (together with the mannequin) stays inside the chosen AWS Area. The mannequin artifacts are imported into the AWS operated deployment account utilizing a digital personal cloud (VPC) endpoint, and you’ll encrypt your mannequin knowledge utilizing an AWS Key Administration Service (AWS KMS) buyer managed key.

Within the following sections, we dive deep into every of those steps to deploy, check, and consider the mannequin.

Stipulations

We use Mistral-7B-v0.3 on this publish as a result of it makes use of an prolonged vocabulary in comparison with its prior model produced by Mistral AI. This mannequin is easy to fine-tune, and Mistral AI has offered instance fine-tuned fashions. We use Mistral for this use case as a result of this mannequin helps a 32,000-token context capability and is fluent in English, French, Italian, German, Spanish, and coding languages. With the Combination of Specialists (MoE) characteristic, it may well obtain greater accuracy for buyer assist use instances.

Mistral-7B-v0.3 is a gated mannequin on the Hugging Face mannequin repository. You should assessment the phrases and situations and request entry to the mannequin by submitting your particulars.

We use Amazon SageMaker Studio to preprocess the information and fine-tune the Mistral mannequin utilizing a SageMaker coaching job. To arrange SageMaker Studio, discuss with Launch Amazon SageMaker Studio. Seek advice from the SageMaker JupyterLab documentation to arrange and launch a JupyterLab pocket book. You’ll submit a SageMaker coaching job to fine-tune the Mistral mannequin from the SageMaker JupyterLab pocket book, which may discovered on the GitHub repo.

Tremendous-tune the mannequin utilizing QLoRA

To fine-tune the Mistral mannequin, we apply QLoRA and Parameter-Environment friendly Tremendous-Tuning (PEFT) optimization strategies. Within the offered pocket book, you employ the Totally Sharded Information Parallel (FSDP) PyTorch API to carry out distributed mannequin tuning. You utilize supervised fine-tuning (SFT) to fine-tune the Mistral mannequin.

Put together the dataset

Step one within the fine-tuning course of is to arrange and format the dataset. After you rework the dataset into the Mistral Default Instruct format, you add it as a JSONL file into the S3 bucket utilized by the SageMaker session, as proven within the following code:

# Load dataset from the hub
dataset = load_dataset("Open-Orca/OpenOrca")
flan_dataset = dataset.filter(lambda instance, indice: "flan" in instance["id"], with_indices=True)
flan_dataset = flan_dataset["train"].train_test_split(test_size=0.01, train_size=0.035)

columns_to_remove = listing(dataset["train"].options)
flan_dataset = flan_dataset.map(create_conversation, remove_columns=columns_to_remove, batched=False)

# save datasets to s3
flan_dataset["train"].to_json(f"{training_input_path}/train_dataset.json", orient="information", force_ascii=False)
flan_dataset["test"].to_json(f"{training_input_path}/test_dataset.json", orient="information", force_ascii=False)

You rework the dataset into Mistral Default Instruct format inside the SageMaker coaching job as instructed within the coaching script (run_fsdp_qlora.py):

    ################
    # Dataset
    ################
    
    train_dataset = load_dataset(
        "json",
        data_files=os.path.be a part of(script_args.dataset_path, "train_dataset.json"),
        cut up="prepare",
    )
    test_dataset = load_dataset(
        "json",
        data_files=os.path.be a part of(script_args.dataset_path, "test_dataset.json"),
        cut up="prepare",
    )

    ################
    # Mannequin & Tokenizer
    ################

    # Tokenizer        
    tokenizer = AutoTokenizer.from_pretrained(script_args.model_id, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.chat_template = MISTRAL_CHAT_TEMPLATE
    
    # template dataset
    def template_dataset(examples):
        return{"textual content":  tokenizer.apply_chat_template(examples["messages"], tokenize=False)}
    
    train_dataset = train_dataset.map(template_dataset, remove_columns=["messages"])
    test_dataset = test_dataset.map(template_dataset, remove_columns=["messages"])

Optimize fine-tuning utilizing QLoRA

You optimize your fine-tuning utilizing QLoRA and with the precision offered as enter into the coaching script as SageMaker coaching job parameters. QLoRA is an environment friendly fine-tuning method that reduces reminiscence utilization to fine-tune a 65-billion-parameter mannequin on a single 48 GB GPU, preserving the complete 16-bit fine-tuning process efficiency. On this pocket book, you employ the bitsandbytes library to arrange quantization configurations, as proven within the following code:

    # Mannequin    
    torch_dtype = torch.bfloat16 if training_args.bf16 else torch.float32
    quant_storage_dtype = torch.bfloat16

    if script_args.use_qlora:
        print(f"Utilizing QLoRA - {torch_dtype}")
        quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch_dtype,
                bnb_4bit_quant_storage=quant_storage_dtype,
            )
    else:
        quantization_config = None

You utilize the LoRA config based mostly on the QLoRA paper and Sebastian Raschka experiment, as proven within the following code. Two key factors to think about from the Raschka experiment are that QLoRA affords 33% reminiscence financial savings at the price of an 39% enhance in runtime, and to verify LoRA is utilized to all layers to maximise mannequin efficiency.

################
# PEFT
################
# LoRA config based mostly on QLoRA paper & Sebastian Raschka experiment
peft_config = LoraConfig(
    lora_alpha=8,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    )

You utilize SFTTrainer to fine-tune the Mistral mannequin:

    ################
    # Coaching
    ################
    coach = SFTTrainer(
        mannequin=mannequin,
        args=training_args,
        train_dataset=train_dataset,
        dataset_text_field="textual content",
        eval_dataset=test_dataset,
        peft_config=peft_config,
        max_seq_length=script_args.max_seq_length,
        tokenizer=tokenizer,
        packing=True,
        dataset_kwargs={
            "add_special_tokens": False,  # We template with particular tokens
            "append_concat_token": False,  # No want so as to add further separator token
        },
    )

On the time of writing, solely merged adapters are supported utilizing the Customized Mannequin Import characteristic for Amazon Bedrock. Let’s have a look at how you can merge the adapter with the bottom mannequin subsequent.

Merge the adapters

Adapters are new modules added between layers of a pre-trained community. Creation of those new modules is feasible by back-propagating gradients via a frozen, 4-bit quantized pre-trained language mannequin into low-rank adapters within the fine-tuning course of. To import the Mistral mannequin into Amazon Bedrock, the adapters must be merged with the bottom mannequin and saved in Safetensors format. Use the next code to merge the mannequin adapters and save them in Safetensors format:

        # load PEFT mannequin in fp16
        mannequin = AutoPeftModelForCausalLM.from_pretrained(
            training_args.output_dir,
            low_cpu_mem_usage=True,
            torch_dtype=torch.float16
        )
        # Merge LoRA and base mannequin and save
        mannequin = mannequin.merge_and_unload()
        mannequin.save_pretrained(
            sagemaker_save_dir, safe_serialization=True, max_shard_size="2GB"
        )

To import the Mistral mannequin into Amazon Bedrock, the mannequin must be in an uncompressed listing inside an S3 bucket accessible by the Amazon Bedrock service position used within the import job.

Import the fine-tuned mannequin into Amazon Bedrock

Now that you’ve got fine-tuned the mannequin, you possibly can import the mannequin into Amazon Bedrock. On this part, we show how you can import the mannequin utilizing the Amazon Bedrock console or the SDK.

Import the mannequin utilizing the Amazon Bedrock console

To import the mannequin utilizing the Amazon Bedrock console, see Import a mannequin with Customized Mannequin Import. You utilize the Import mannequin web page as proven within the following screenshot to import the mannequin from the S3 bucket.

After you efficiently import the fine-tuned mannequin, you possibly can see the mannequin listed on the Amazon Bedrock console.

Import the mannequin utilizing the SDK

The AWS Boto3 library helps importing customized fashions into Amazon Bedrock. You should use the next code to import a fine-tuned mannequin from inside the pocket book into Amazon Bedrock. That is an asynchronous methodology.

import boto3
import datetime
br_client = boto3.consumer('bedrock', region_name="")
pt_model_nm = ""
pt_imp_jb_nm = f"{pt_model_nm}-{datetime.datetime.now().strftime('%YpercentmpercentdpercentMpercentHpercentS')}"
role_arn = "<>"
pt_model_src = {"s3DataSource": {"s3Uri": f"{pt_pubmed_model_s3_path}"}}
resp = br_client.create_model_import_job(jobName=pt_imp_jb_nm,
                                  importedModelName=pt_model_nm,
                                  roleArn=role_arn,
                                  modelDataSource=pt_model_src)

Take a look at the imported mannequin

Now that you’ve got imported the fine-tuned mannequin into Amazon Bedrock, you possibly can check the mannequin. On this part, we show how you can check the mannequin utilizing the Amazon Bedrock console or the SDK.

Take a look at the mannequin on the Amazon Bedrock console

You’ll be able to check the imported mannequin utilizing an Amazon Bedrock playground, as illustrated within the following screenshot.

Take a look at the mannequin utilizing the SDK

You can even use the Amazon Bedrock Invoke Mannequin API to run the fine-tuned imported mannequin, as proven within the following code:

consumer = boto3.consumer("bedrock-runtime", region_name="us-west-2")
model_id = "<>"


def call_invoke_model_and_print(native_request):
    request = json.dumps(native_request)

    strive:
        # Invoke the mannequin with the request.
        response = consumer.invoke_model(modelId=model_id, physique=request)
        model_response = json.hundreds(response["body"].learn())

        response_text = model_response["outputs"][0]["text"]
        print(response_text)
    besides (ClientError, Exception) as e:
        print(f"ERROR: Cannot invoke '{model_id}'. Purpose: {e}")
        exit(1)

immediate = "will there be a season 5 of shadowhunters"
formatted_prompt = f"[INST] {immediate} [/INST]"
native_request = {
"immediate": formatted_prompt,
"max_tokens": 64,
"top_p": 0.9,
"temperature": 0.91
}
call_invoke_model_and_print(native_request)

The customized Mistral mannequin that you simply imported utilizing Amazon Bedrock helps temperature, top_p, and max_gen_len parameters when invoking the mannequin for inferencing. The inference parameters top_k, max_seq_len, max_batch_size, and max_new_tokens should not supported for a customized Mistral fine-tuned mannequin.

Consider the imported mannequin

Now that you’ve got imported and examined the mannequin, let’s consider the imported mannequin utilizing the SageMaker FMEval library. For extra particulars, discuss with Consider Bedrock Imported Fashions. To judge the query answering process, we use the metrics F1 Rating, Actual Match Rating, Quasi Actual Match Rating, Precision Over Phrases, and Recall Over Phrases. The important thing metrics for the query answering duties are Actual Match, Quasi-Actual Match, and F1 over phrases evaluated by evaluating the mannequin predicted solutions towards the bottom fact solutions. The FMEval library helps out-of-the-box analysis algorithms for metrics reminiscent of accuracy, QA Accuracy, and others detailed within the FMEval documentation. Since you fine-tuned the Mistral mannequin for query answering, you should use the QA Accuracy algorithm, as proven within the following code. The FMEval library helps these metrics for the QA Accuracy algorithm.

config = DataConfig(
    dataset_name="trex_sample",
    dataset_uri="knowledge/test_dataset.json",
    dataset_mime_type=MIME_TYPE_JSONLINES,
    model_input_location="query",
    target_output_location="reply"
)
bedrock_model_runner = BedrockModelRunner(
    model_id=model_id,
    output="outputs[0].textual content",
    content_template="{"immediate": $immediate, "max_tokens": 500}",
)

eval_algo = QAAccuracy()
eval_output = eval_algo.consider(mannequin=bedrock_model_runner, dataset_config=config, 
                                    prompt_template="[INST]$model_input[/INST]", save=True)

You will get the consolidated metrics for the imported mannequin as follows:

for op in eval_output:
    print(f"Eval Title: {op.eval_name}")
    for rating in op.dataset_scores:
        print(f"{rating.title} : {rating.worth}")

Clear up

To delete the imported mannequin from Amazon Bedrock, navigate to the mannequin on the Amazon Bedrock console. On the choices menu (three dots), select Delete.

To delete the SageMaker area together with the SageMaker JupyterLab house, discuss with Delete an Amazon SageMaker area. You might also need to delete the S3 buckets the place the information and mannequin are saved. For directions, see Deleting a bucket.

Conclusion

On this publish, we defined the totally different facets of fine-tuning a Mistral mannequin utilizing SageMaker, importing the mannequin into Amazon Bedrock, invoking the mannequin utilizing each an Amazon Bedrock playground and Boto3, after which evaluating the imported mannequin utilizing the FMEval library. You should use this characteristic to import base FMs or FMs fine-tuned both on premises, on SageMaker, or on Amazon EC2 into Amazon Bedrock and use the fashions with none heavy lifting in your generative AI purposes. Discover the Customized Mannequin Import characteristic for Amazon Bedrock to deploy FMs fine-tuned for code technology duties in a safe and scalable method. Go to our GitHub repository to discover samples ready for fine-tuning and importing fashions from numerous households.

Concerning the Authors

Jay Pillai is a Principal Options Architect at Amazon Net Companies. On this position, he capabilities because the Lead Architect, serving to companions ideate, construct, and launch Companion Options. As an Data Know-how Chief, Jay makes a speciality of synthetic intelligence, generative AI, knowledge integration, enterprise intelligence, and person interface domains. He holds 23 years of intensive expertise working with a number of shoppers throughout provide chain, authorized applied sciences, actual property, monetary providers, insurance coverage, funds, and market analysis enterprise domains.

Rupinder Grewal is a Senior AI/ML Specialist Options Architect with AWS. He at present focuses on serving of fashions and MLOps on Amazon SageMaker. Previous to this position, he labored as a Machine Studying Engineer constructing and internet hosting fashions. Outdoors of labor, he enjoys enjoying tennis and biking on mountain trails.

Evandro Franco is a Sr. AI/ML Specialist Options Architect at Amazon Net Companies. He helps AWS clients overcome enterprise challenges associated to AI/ML on high of AWS. He has greater than 18 years of expertise working with expertise, from software program improvement, infrastructure, serverless, to machine studying.

Felipe Lopez is a Senior AI/ML Specialist Options Architect at AWS. Previous to becoming a member of AWS, Felipe labored with GE Digital and SLB, the place he targeted on modeling and optimization merchandise for industrial purposes.

Sandeep Singh is a Senior Generative AI Information Scientist at Amazon Net Companies, serving to companies innovate with generative AI. He makes a speciality of generative AI, synthetic intelligence, machine studying, and system design. He’s keen about growing state-of-the-art AI/ML-powered options to unravel complicated enterprise issues for various industries, optimizing effectivity and scalability.

Ragha Prasad is a Principal Engineer and a founding member of Amazon Bedrock, the place he has had the privilege to take heed to buyer wants first-hand and understands what it takes to construct and launch scalable and safe Gen AI merchandise. Previous to Bedrock, he labored on quite a few merchandise in Amazon, starting from units to Adverts to Robotics.

Paras Mehra is a Senior Product Supervisor at AWS. He’s targeted on serving to construct Amazon SageMaker Coaching and Processing. In his spare time, Paras enjoys spending time along with his household and street biking across the Bay Space.