Today we are pleased to announce that the DBRX modelan open, general-purpose large language model (LLM) developed by Data bricks, is available for customers through amazon SageMaker JumpStart to deploy with a single click to run inference. The DBRX LLM employs a detailed Mixture of Experts (MoE) architecture, pre-trained on 12 billion carefully selected data tokens 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 to get you started quickly with ML. In this post, we explain how to discover and implement the DBRX model.
What is the DBRX model?
DBRX is a sophisticated set-top-only LLM built on a transformer architecture. It employs a detailed MoE architecture, incorporating 132 billion total parameters, of which 36 billion are active for any given input.
The model was pre-trained using a data set consisting of 12 billion text and code tokens. Unlike other open MoE models such as Mixtral and Grok-1, DBRX features a fine-grained approach, using a larger number of smaller experts for optimized performance. Compared with other MoE models, DBRX has 16 experts and chooses 4.
The model is available under the Databricks Open Model license, for unrestricted use.
What is SageMaker JumpStart?
SageMaker JumpStart is a fully managed platform that offers next-generation core models for a variety of use cases, including content writing, code generation, question answering, writing, summarizing, classification, and information retrieval. It provides a collection of pre-trained models that you can deploy quickly and easily, accelerating the development and deployment of machine learning applications. One of the key components of SageMaker JumpStart is Model Hub, which offers a wide catalog of pre-trained models, such as DBRX, for a variety of tasks.
You can now discover and deploy DBRX models with a few clicks in amazon SageMaker Studio or programmatically through the SageMaker Python SDK, allowing you to get model performance and MLOps controls with amazon SageMaker features such as amazon SageMaker Pipelines, amazon SageMaker Debugger or container logs. . The model is deployed in a secure AWS environment and under the controls of your VPC, which helps provide data security.
Discover models 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 go over 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 tools specifically designed to perform all ML development steps, from data preparation to authoring, training, and deployment. implementation of your ML models. For more details about getting started and setting up SageMaker Studio, see amazon SageMaker Studio.
In SageMaker Studio, you can access SageMaker JumpStart by choosing Good start in the navigation panel.
From the SageMaker JumpStart home page, you can search for “DBRX” in the search box. The search results will show DBRX Instruction and DBRX base.
You can choose the model card to view details about the model, such as the license, the data used to train, and how to use the model. You will also find the Deploy to deploy the model and create an endpoint.
Deploy the model to SageMaker JumpStart
Deployment begins when you choose the Deploy button. Once the deployment is complete, you will see an endpoint being 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'll see sample code that you can use in the notebook editor of your choice in SageMaker Studio.
DBRX base
To deploy using the SDK, we start by selecting the base DBRX model, specified by model_id with value huggingface-llm-dbrx-base. You can deploy any of the selected models in SageMaker with the following code. Similarly, you can implement DBRX Instruct using your own model 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 model to SageMaker with default configurations, including the default instance type and default VPC configurations. You can change these settings by specifying non-default values in JumpStartModel. 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 service limit at the account level to use ml.p4d.24xlarge or ml.pde.24xlarge for endpoint use as one or more instances. You can follow the instructions here to request a service fee increase.
Once deployed, you can run inference against the deployed endpoint via the SageMaker predictor:
payload = {
"inputs": "Hello!",
"parameters": {
"max_new_tokens": 10,
},
}
predictor.predict(payload)
Example 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 messages and sample results.
GENERATION CODE
Using the above example, we can use code generation messages 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)
The following 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
Analysis of feelings
You can perform sentiment analysis using a message like the following 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)
The following is the result:
Answer to questions
You can use a question response message like the following 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)
The following 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 version adapted to DBRX instructions accepts formatted instructions in which conversation roles must begin with a message from the user and alternate between user and assistant instructions (DBRX instructions). The instructional format must be strictly adhered to; otherwise the model will generate suboptimal results. The template for creating a message for the Instruct model 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|> They are special tokens for the beginning of the chain (BOS) and the end of the chain (EOS). The model can contain multiple turns of conversation between the system, the user, and the assistant, allowing the incorporation of short examples to improve the model's responses.
The following code shows how you can format the message 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 message to retrieve knowledge:
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)
The following 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.
GENERATION CODE
DBRX models demonstrate comparative 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)
The following 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 report strengths in mathematical precision. 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 output 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.
Clean
Once you have finished running the notebook, be sure to delete all the resources you created in the process so that your billing stops. Use the following code:
predictor.delete_model()
predictor.delete_endpoint()
Conclusion
In this post, we show you how to get started with DBRX in SageMaker Studio and deploy the model for inference. Since the base models are pre-trained, they can help reduce training and infrastructure costs and allow customization for 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 and works with a leading global systems integrator. He earned the title of one of the youngest Indian master inventors with over 400 patents in the domains of ai/ML and IoT. He has 8+ years of industry experience, from startups to large-scale enterprises, from IoT research engineer, data scientist to data and ai architect. Shikhar assists in architecting, building and maintaining scalable and cost-effective cloud environments for organizations and supports GSI partners in building strategic industries.
Nithiyn Vijeaswaran He is a solutions architect at AWS. His area of focus is Generative ai and AWS ai Accelerators. He has 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 an avid Dallas Mavericks fan and enjoys collecting sneakers.
Sebastian Bustillo He is a solutions architect at AWS. He focuses on ai/ML technologies with a deep passion for generative ai and computing accelerators. At AWS, he helps customers unlock business value through generative ai. When he's not at work, he likes to brew a perfect cup of specialty coffee and explore the world with his wife.
Armando Diaz He is a solutions architect at AWS. He 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 he is not at work, he enjoys spending time with his wife and his family, hiking, and traveling the world.
