Today we are delighted to announce that the Model DBRXan open, general-purpose large language model (LLM) developed by Data bricksis available to customers via Amazon SageMaker JumpStart to deploy with a single click to run inference. The DBRX LLM uses a fine-grained mixture of experts (MoE) architecture, pre-trained on 12 trillion tokens of carefully curated data and a maximum context length of 32,000 tokens.
You can try this model with SageMaker JumpStart, a machine learning (ML) hub that provides access to algorithms and models so you can quickly get started with ML. In this article, we explain how to discover and deploy the DBRX model.
What is the DBRX model
DBRX is a sophisticated set-top box-only LLM, built on a transformer architecture. It uses a fine-grained MoE architecture, incorporating 132 billion parameters in total, of which 36 billion are active for a given input.
The model underwent pre-training using a dataset consisting of 12 trillion text and code tokens. Unlike other open MoE models like Mixtral and Grok-1, DBRX offers a more fine-grained approach, using a larger number of smaller experts for optimized performance. Compared to other MoE models, DBRX has 16 experts and chooses 4.
The model is made available under the Databricks Open Model license, for unrestricted use.
What is SageMaker JumpStart?
SageMaker JumpStart is a fully managed platform that offers industry-leading core templates for various use cases such as content writing, code generation, question answering, drafting, summarizing, classifying, and information retrieval. It provides a collection of pre-trained models that you can deploy quickly and easily, accelerating the development and deployment of ML applications. One of the key components of SageMaker JumpStart is the Model Hub, which offers an extensive catalog of pre-trained models, such as DBRX, for a variety of tasks.
You can now discover and deploy DBRX models in just a few clicks Amazon SageMaker Studio or programmatically via the SageMaker Python SDK, allowing you to derive model performance and MLOps controls with Amazon SageMaker features such as Amazon SageMaker Pipelines, Amazon SageMaker Debuggeror container logs. The model is deployed in a secure AWS environment and under the controls of your VPC, helping to ensure data security.
Explore templates in SageMaker JumpStart
You can access the DBRX model through SageMaker JumpStart in the SageMaker Studio UI and the SageMaker Python SDK. In this section, we explain how to discover models in SageMaker Studio.
SageMaker Studio is an integrated development environment (IDE) that provides a single web-based visual interface where you can access purpose-built tools to perform all stages of ML development, from data preparation to authoring, training, and deployment of your ML models. For details on how to start and configure SageMaker Studio, see Amazon SageMaker Studio.
In SageMaker Studio, you can access SageMaker JumpStart by choosing Start of jump in the navigation pane.
![](https://d2908q01vomqb2.cloudfront.net/f1f836cb4ea6efb2a0b1b99f41ad8b103eff4b59/2024/04/24/DBRX001.jpg)
From the SageMaker JumpStart landing page, you can search for “DBRX” in the search box. The search results will list DBRX instruction And DBRX database.
![](https://d2908q01vomqb2.cloudfront.net/f1f836cb4ea6efb2a0b1b99f41ad8b103eff4b59/2024/04/24/DBRX003.jpg)
You can choose the model card to view model details, such as license, data used for training, and how to use the model. You will also find the Deploy to deploy the template and create an endpoint.
![](https://d2908q01vomqb2.cloudfront.net/f1f836cb4ea6efb2a0b1b99f41ad8b103eff4b59/2024/04/24/dbrx005.jpg)
Deploy the model in SageMaker JumpStart
Deployment begins when you choose the Deploy button. Once the deployment is complete, you will see that an endpoint is created. You can test the endpoint by passing a sample inference request payload or by selecting the test option using the SDK. When you select the option to use the SDK, you will see sample code that you can use in your notebook editor of choice in SageMaker Studio.
DBRX database
To deploy using the SDK, we start by selecting the DBRX Base template, specified by the model_id
with the value huggingface-llm-dbrx-base. You can deploy any of the selected models to SageMaker with the following code. Similarly, you can deploy DBRX Instruct using its own template ID.
from sagemaker.jumpstart.model import JumpStartModel
accept_eula = True
model = JumpStartModel(model_id="huggingface-llm-dbrx-base")
predictor = model.deploy(accept_eula=accept_eula)
This deploys the template to SageMaker with default configurations, including the default instance type and default VPC configurations. You can modify these configurations by specifying non-default values in JumpStart Model. The Eula value must be explicitly set to True in order to accept the End User License Agreement (EULA). Also make sure you have the account-level service limit to use ml.p4d.24xlarge or ml.pde.24xlarge for endpoint use as one or more instances. You can follow the instructions here in order to request an increase in the service quota.
Once deployed, you can run inference on the deployed endpoint via the SageMaker Predictor:
payload = {
"inputs": "Hello!",
"parameters": {
"max_new_tokens": 10,
},
}
predictor.predict(payload)
Sample Prompts
You can interact with the DBRX base model like any standard text generation model, where the model processes an input sequence and generates the next predicted words in the sequence. In this section, we provide some example prompts and results.
Code generation
Using the previous example, we can use the code generation prompts as follows:
payload = {
"inputs": "Write a function to read a CSV file in Python using pandas library:",
"parameters": {
"max_new_tokens": 30, }, }
response = predictor.predict(payload)("generated_text").strip()
print(response)
Here is the result :
import pandas as pd
df = pd.read_csv("file_name.csv")
#The above code will import pandas library and then read the CSV file using read_csv
Sentiment analysis
You can perform sentiment analysis using a prompt like this with DBRX:
payload = {
"inputs": """
Tweet: "I am so excited for the weekend!"
Sentiment: Positive
Tweet: "Why does traffic have to be so terrible?"
Sentiment: Negative
Tweet: "Just saw a great movie, would recommend it."
Sentiment: Positive
Tweet: "According to the weather report, it will be cloudy today."
Sentiment: Neutral
Tweet: "This restaurant is absolutely terrible."
Sentiment: Negative
Tweet: "I love spending time with my family."
Sentiment:""",
"parameters": {
"max_new_tokens": 2,
},
}
response = predictor.predict(payload)("generated_text").strip()
print(response)
Here is the result :
Answer to questions
You can use a question answer prompt like this with DBRX:
# Question answering
payload = {
"inputs": "Respond to the question: How did the development of transportation systems, such as railroads and steamships, impact global trade and cultural exchange?",
"parameters": {
"max_new_tokens": 225,
},
}
response = predictor.predict(payload)("generated_text").strip()
print(response)
Here is the result :
The development of transportation systems, such as railroads and steamships, impacted global trade and cultural exchange in a number of ways.
The documents provided show that the development of these systems had a profound effect on the way people and goods were able to move around the world.
One of the most significant impacts of the development of transportation systems was the way it facilitated global trade.
The documents show that the development of railroads and steamships made it possible for goods to be transported more quickly and efficiently than ever before.
This allowed for a greater exchange of goods between different parts of the world, which in turn led to a greater exchange of ideas and cultures.
Another impact of the development of transportation systems was the way it facilitated cultural exchange. The documents show that the development of railroads and steamships made it possible for people to travel more easily and quickly than ever before.
This allowed for a greater exchange of ideas and cultures between different parts of the world. Overall, the development of transportation systems, such as railroads and steamships, had a profound impact on global trade and cultural exchange.
DBRX instruction
The instruction-optimized version of DBRX accepts formatted instructions in which chat roles must begin with a user prompt and alternate between user and wizard instructions (DBRX-instruct). The format of the instructions must be strictly adhered to, otherwise the model will generate suboptimal outputs. The template for creating a prompt for the Instruct template is defined as follows:
<|im_start|>system
{system_message} <|im_end|>
<|im_start|>user
{human_message} <|im_end|>
<|im_start|>assistant\n
<|im_start|>
And <|im_end|>
are special tokens for start of chain (BOS) and end of chain (EOS). The model can contain several turns of conversation between the system, user and assistant, allowing the incorporation of a few examples to improve the model's responses.
The following code shows how to format the prompt in statement format:
from typing import Dict, List
def format_instructions(instructions: List(Dict(str, str))) -> List(str):
"""Format instructions where conversation roles must alternate system/user/assistant/user/assistant/..."""
prompt: List(str) = ()
for instruction in instructions:
if instruction("role") == "system":
prompt.extend(("<|im_start|>system\n", (instruction("content")).strip(), " <|im_end|>\n"))
elif instruction("role") == "user":
prompt.extend(("<|im_start|>user\n", (instruction("content")).strip(), " <|im_end|>\n"))
else:
raise ValueError(f"Invalid role: {instruction('role')}. Role must be either 'user' or 'system'.")
prompt.extend(("<|im_start|>assistant\n"))
return "".join(prompt)
def print_instructions(prompt: str, response: str) -> None:
bold, unbold = '\033(1m', '\033(0m'
print(f"{bold}> Input{unbold}\n{prompt}\n\n{bold}> Output{unbold}\n{response('generated_text').strip()}\n")
Knowledge recovery
You can use the following prompt for knowledge retrieval:
instructions = (
{"role": "system", "content": "You are an expert in suggesting diet plans optimized for different sports or fitness activities, taking into account energy requirements, and recovery requirements."},
{"role": "user", "content": "I am trying to build muscle mass and increase my strength. What kind of high-protein, calorie-dense diet plan would you recommend for me?"}
)
prompt = format_instructions(instructions)
payload = {
"inputs": prompt,
"parameters": {"max_new_tokens": 700, "temperature": 0.5, "do_sample": True}
}
response=predictor.predict(payload)
print_instructions(prompt, response)
Here is the result :
> Input
<|im_start|>system
You are an expert in suggesting diet plans optimized for different sports or fitness activities, taking into account energy requirements, and recovery requirements. <|im_end|>
<|im_start|>user
I am trying to build muscle mass and increase my strength. What kind of high-protein, calorie-dense diet plan would you recommend for me? <|im_end|>
<|im_start|>assistant
> Output
Absolutely, I'd be happy to help you with that. Given your goal of building muscle mass and increasing strength, you'll want a diet that's high in protein and provides sufficient calories to support muscle growth and recovery. Here's a basic outline of what your diet plan could look like:
1. **Breakfast:**
- 3 whole eggs
- 3 egg whites
- 1 cup of oatmeal
- 1 scoop of whey protein
- 1 banana
2. **Mid-Morning Snack:**
- 1.5 ounces of almonds
- 1 apple
3. **Lunch:**
- 8 ounces of chicken breast
- 1 cup of brown rice
- 1 cup of mixed vegetables
4. **Afternoon Snack:**
- 1.5 cups of Greek yogurt
- 1 cup of mixed berries
5. **Dinner:**
- 8 ounces of salmon or lean beef
- 1 cup of quinoa
- 1.5 cups of mixed vegetables
6. **Evening Snack:**
- 1 scoop of casein protein
- 1 tablespoon of almond butter
This is a high-protein, calorie-dense diet that includes a variety of nutrient-rich foods. It's important to note that everyone's nutritional needs are different, so you may need to adjust portion sizes or food choices based on your specific needs and preferences.
Also, remember to stay hydrated by drinking plenty of water throughout the day. Regular exercise, particularly strength training, is crucial to achieving your goals. It's also important to get enough sleep to aid in muscle recovery and growth.
Before starting any new diet or exercise program, it's a good idea to consult with a healthcare professional or a registered dietitian. They can provide personalized advice based on your health history and current lifestyle.
Code generation
DBRX models demonstrate proven strengths for coding tasks. For example, see the following code:
instructions = (
{"role": "system", "content": "You are an expert in coding and programming. Your role is to provide examples, explanations, and insights related to programming languages, algorithms, data structures, and software development practices."},
{"role": "user", "content": "Write a Python script to read a CSV file containing stock prices and plot the closing prices over time using Matplotlib."}
)
prompt = format_instructions(instructions)
payload = { "inputs": prompt, "parameters": {"max_new_tokens": 315, "temperature": 0.3, "do_sample": True} }
response=predictor.predict(payload)
print_instructions(prompt, response)
Here is the result :
> Input
<|im_start|>system
You are an expert in coding and programming. Your role is to provide examples, explanations, and insights related to programming languages, algorithms, data structures, and software development practices. <|im_end|>
<|im_start|>user
Write a Python script to read a CSV file containing stock prices and plot the closing prices over time using Matplotlib. <|im_end|>
<|im_start|>assistant
> Output
Sure, I'd be happy to help you with that. Here's a simple Python script that uses the `pandas` library to read a CSV file and `matplotlib` to plot the closing prices over time.
```python
import pandas as pd
import matplotlib.pyplot as plt
# Load the data from a CSV file
df = pd.read_csv('stock_prices.csv')
# Plot the closing prices over time
plt.figure(figsize=(12,5))
plt.plot(df('Date'), df('Close'))
plt.xlabel('Date')
plt.ylabel('Price')
plt.title('Stock Prices Over Time')
plt.show()
```
In this script, replace `'stock_prices.csv'` with the path to your CSV file. The file should have columns named 'Date' and 'Close' for this script to work correctly. If your file uses different column names, you'll need to adjust the script accordingly.
The `pandas` library is a powerful tool for data manipulation and analysis. It can read data from a variety of sources, including CSV files, and represent the data in a tabular format. The `matplotlib` library is a popular choice for creating static, animated, and interactive visualizations in Python.
This script first loads the data from the CSV file into a pandas DataFrame. It then plots the 'Close' column against the 'Date' column using matplotlib's `plot()` function. The `figure()` function is used to specify the size of the plot, and `show()` is used to display the plot.
Mathematics and reasoning
DBRX models also demonstrate strengths in mathematical accuracy. For example, see the following code:
instructions = (
{"role": "system", "content": "You are an expert in mathematics and reasoning. Your role is to provide examples, explanations, and insights related to mathematical concepts, problem-solving techniques, and logical reasoning."},
{"role": "user", "content": "A car rental company charges $40 per day plus $0.25 per mile driven. If I rented a car for 3 days and drove 250 miles, what would be the total cost? Break down the calculation steps before giving the final answer."}
)
prompt = format_instructions(instructions)
payload = {
"inputs": prompt,
"parameters": {"max_new_tokens": 400, "temperature": 0.2, "do_sample": True}
}
response=predictor.predict(payload)
print_instructions(prompt, response)
DBRX can provide understanding as shown in the following result with mathematical logic:
> Input
<|im_start|>system
You are an expert in mathematics and reasoning. Your role is to provide examples, explanations, and insights related to mathematical concepts, problem-solving techniques, and logical reasoning. <|im_end|>
<|im_start|>user
A car rental company charges $40 per day plus $0.25 per mile driven. If I rented a car for 3 days and drove 250 miles, what would be the total cost? Break down the calculation steps before giving the final answer. <|im_end|>
<|im_start|>assistant
> Output
Sure, I'd be happy to help with that! The problem you've presented involves both fixed costs (the daily rate) and variable costs (the cost per mile). Let's break it down:
1. First, let's calculate the total cost of the fixed daily rate. The car rental company charges $40 per day, and you rented the car for 3 days. So, the total cost for the daily rate is $40/day * 3 days = $120.
2. Next, let's calculate the total cost of the variable rate, which is the cost per mile driven. The company charges $0.25 per mile, and you drove 250 miles. So, the total cost for the mileage is $0.25/mile * 250 miles = $62.50.
3. Finally, to find the total cost of renting the car, we add the total cost of the daily rate and the total cost of the mileage. So, the total cost is $120 + $62.50 = $182.50.
So, the total cost of renting the car for 3 days and driving 250 miles would be $182.50.
To clean
Once you have finished running the notebook, be sure to delete any resources you created in the process so that your billing is stopped. Use the following code:
predictor.delete_model()
predictor.delete_endpoint()
Conclusion
In this article, we showed you how to get started with DBRX in SageMaker Studio and deploy the model for inference. Because base models are pre-trained, they can help reduce training and infrastructure costs and enable customization to suit your use case. Visit SageMaker JumpStart in SageMaker Studio now to get started.
Resources
about the authors
Shikhar Kwatra is a Solutions Architect specializing in AI/ML at Amazon Web Services, working with a leading global systems integrator. He has earned the title of one of India's youngest master inventors with over 400 patents in the fields of AI/ML and IoT. He has over 8 years of industry experience, from startups to large enterprises, from IoT Research Engineer, Data Scientist, to Data & AI Architect. Shikhar helps design, build and maintain cost-effective and scalable cloud environments for organizations and supports GSI partners in building a strategic industry.
Niithiyn Vijeaswaran is a solutions architect at AWS. His area of focus is generative AI and AWS AI accelerators. He holds a bachelor's degree in computer science and bioinformatics. Niithiyn works closely with the Generative AI GTM team to help AWS customers on multiple fronts and accelerate their adoption of generative AI. He is a huge Dallas Mavericks fan and enjoys collecting sneakers.
Sébastien Bustillo is a solutions architect at AWS. He focuses on AI/ML technologies with a deep passion for generative AI and computational accelerators. At AWS, he helps customers unlock business value through generative AI. When he's not at work, he enjoys brewing a perfect cup of specialty coffee and exploring the world with his wife.
Armando Diaz is a solutions architect at AWS. It focuses on generative AI, AI/ML and data analytics. At AWS, Armando helps customers integrate cutting-edge generative AI capabilities into their systems, driving innovation and competitive advantage. When not at work, he enjoys spending time with his wife and family, hiking, and traveling the world.