Use AI Builder: Model specialization to create a gen AI agent for text
This feature is in Beta.
This article describes how to create a generative AI agent for custom text-based tasks using AI Builder: Model specialization.
AI Builder provides a simple, no-code approach to build and optimize domain-specific, high-quality AI agent systems for common AI use cases.
What can you do with Model specialization?
Use AI Builder: Model specialization to generate high-quality results for any domain-specific task, such as summarization, classification, text transformation, and content generation.
Examples of model specialization include:
- Summarizing the issue and resolution of customer calls.
- Analyzing the sentiment of customer reviews.
- Classifying research papers by topic.
- Generating press releases for new features.
Given high-level instruction and examples, AI Builder: Model specialization optimizes prompts on behalf of users, automatically infers evaluation criteria, evaluates the system from provided data, and deploys the model as a productionizable endpoint.
AI Builder: Model specialization leverages automated evaluation capabilities, including MLflow and Agent Evaluation, to enable rapid assessment of the cost-quality tradeoff for your specific extraction task. This assessment allows you to make informed decisions about the balance between accuracy and resource investment.
Requirements
- Serverless-supported workspace that includes the following:
- Unity Catalog enabled on your workspace.
- Partner-powered AI assistive features enabled on your workspace.
- A workspace in one of the supported regions:
us-east-1
orus-west-2
. - Access to Mosaic AI Model Serving.
- Access to foundation models in Unity Catalog through the
system.ai
schema. - Access to a serverless budget policy with a nonzero budget.
- Ability to use the
ai_query
SQL function. - You must have input data ready to use. You can choose to provide either:
- A Unity Catalog table. The table name cannot contain any special characters (such as
-
). - At least 3 example inputs and outputs. If you choose this option, you'll need to specify a Unity Catalog schema destination path for the agent, and you must have CREATE REGISTERED MODEL and CREATE TABLE permissions to this schema.
- A Unity Catalog table. The table name cannot contain any special characters (such as
- If you want to optimize your agent, you need at least 100 inputs (either 100 rows in a Unity Catalog table or 100 manually-provided examples).
Create a model specialization agent
Go to AI Builder in the left navigation pane of your workspace and click Model Specialization.
Step 1: Configure your agent
On the Configure tab, click Show an example > to expand an example input and model response for a model specialization agent.
In the pane below, configure your agent:
-
Under Describe your task, enter a clear and detailed description of your specialization task, including its purpose and desired outcome.
-
Under Do you have data?, select Yes or No.
The following data types are supported:
string
,int
, anddouble
.If you select Yes:
-
Under Select Dataset, click Browse to select the table in Unity Catalog you want to use. This dataset is used to create your agent and can be labeled or unlabeled. The table name cannot contain any special characters (such as
-
).The following is an example:
main.model_specialization.customer_call_transcripts
-
In the Input column field, select the column you want to use as your input text. The dropdown menu is automatically populated with columns from your selected table.
-
(Optional) In the Output column (optional), select the column you want to provide as an example output for the expected transformation. Providing this data helps configure your agent to more accurately adapt to your domain-specific needs.
If you select No:
- Provide at least 3 examples of inputs and expected outputs for your specialization task. Providing high-quality examples helps configure your specialization agent to better understand your requirements.
- To add more examples, click + Add.
- Under Agent destination, select the Unity Catalog schema where you’d like AI Builder to help you create a table with evaluation data. You must have CREATE REGISTERED MODEL and CREATE TABLE permissions to this schema.
-
-
Name your agent.
-
Click Create agent.
Step 2: Build and improve your agent
In the Build tab, you can review recommendations to improve your agent, review sample model outputs, and adjust your task instructions and evaluation criteria.
In the Recommendation pane, Databricks provides recommendations to help you define evaluation metrics for your agent and evaluate sample responses as good or bad.
-
Review the Databricks recommendations for optimizing agent performance.
-
Review suggested evaluation criteria. These recommended evaluation criteria are automatically inferred to help you optimize your agent.
For each recommendation:
- To accept the recommendation, select Yes. This adds the evaluation criteria in the Agent configuration pane.
- To reject the criteria, select No.
- You can also choose to Dismiss the recommendation.
-
Review sample model inputs and outputs and provide optional human feedback. This evaluation helps improve the model’s reponses.
For each sample, select whether or not it was a good response. If No, provide optional feedback on the response and click Save to move onto the next one.
-
After you’ve finished reviewing recommendations, review the Agent configuration pane.
- You can adjust the task instructions to be more specific to improve the model’s performance.
- Review the evaluation criteria you added from the recommendations. You can remove criteria by clicking X.
- If you’d like to add more evaluation criteria, click + Add to add your own.
-
Click Update agent to save those changes to your agent. The examples under Review model output update to show new example model outputs.
Step 3: Try out and optimize your agent
Try out your agent in workflows across Databricks.
On the Use tab,
-
Click Try in SQL to open the SQL editor and use
ai_query
to send requests to your new model specialization agent. -
(Optional) Click Optimize if you want to optimize your agent for cost.
- Optimization requires at least 100 inputs. If you provided a Unity Catalog dataset, the table must contain at least 100 rows. If you did not provide a dataset, you need to provide at least 100 examples.
- Optimization can take about an hour.
- Making changes to your currently active agent is blocked when optimization is in progress.
When optimization completes, you are directed to the Review tab to view a comparison of your currently active agent and an agent optimized for cost. See (Optional) Step 4: Review and deploy an optimized agent.
(Optional) Step 4: Review and deploy an optimized agent
Databricks recommends at least 100 inputs (either 100 rows in your Unity Catalog table or 100 manually-provided examples) to optimize your agent. When you add more inputs, the knowledge base that the agent can learn from increases, which improves agent quality and its response accuracy.
When you select Optimize on the Use tab, Databricks compares multiple different optimization strategies to build and deploy an optimized agent. These strategies include Foundation Model Fine-tuning which uses Databricks Geos.
On the Review tab,
- In Evaluation results, you can review the evaluation metrics for the optimized agent. To perform evaluation, Databricks uses metrics based on the evaluation criteria you defined in the Build tab
- Click on a request to open more details. Here, you can see a detailed assessment of each evaluation metric, including the rationale for why it passed or failed. This uses Databricks built-in AI judges. You can also inspect the input and response.
- After you review these results, you can use your optimized agent.
Limitations
- Databricks recommends at least 100 inputs (either 100 rows in your Unity Catalog table or 100 manually-provided samples) to optimize your agent. When you add more inputs, the knowledge base that the agent can learn from increases, which improves agent quality and its response accuracy.
- If you provide a Unity Catalog table, the table name cannot contain any special characters (such as
-
). - Only the following data types are supported as inputs:
string
,int
, anddouble
. - Usage capacity is currently limited to 100k input and output tokens per minute.
- Workspaces that use PrivateLink, including storage behind PrivateLink, are not supported.