Generating database query using machine-learned large language models
A computer system uses a machine-learned language model to generate an SQL query for a user query. The system receives a user query comprising a task for performing a database query. The system identifies an embedding for the user query to represent the user query. The system generates a prompt for input to a machine-learned language model, and the prompt specifies the user query, metadata associated with the identified data table and a request to generate one or more SQL statements for performing the database query on the data table. The system provides the prompt to a model serving system and receives an output generated that includes the requested SQL statements for performing the database query. The system presents a response to the user query using the received SQL statements.
1 . A method comprising:
receiving, from a client device, a user query from a user, the user query comprising a task for performing a database query;
identifying an embedding for the user query to represent the user query in a latent space;
generating a prompt for input to a machine-learned language model, the prompt specifying the user query and a request to generate one or more SQL statements for performing the database query on a data table, wherein the machine-learned language model is trained by:
identifying one or more SQL dashboards implemented for the online system, wherein a SQL dashboard is powered by at least one SQL query,
for each identified SQL dashboard, obtaining a training example including a pair of inputs including metadata describing the SQL dashboard, and outputs including the at least one SQL query powering the identified SQL dashboard, the metadata including one or a combination of column features of data tables associated with the SQL dashboard, headings, descriptions, user comments, or text input for the SQL dashboard,
obtaining a loss indicating a difference between the outputs of the training example and outputs obtained by applying the machine-learned language model to the inputs of the training example, and
fine-tuning parameters of the machine-learned model to reduce the loss based on the metadata obtained for the one or more SQL dashboards and the at least one SQL query powering the one or more SQL dashboards;
providing the prompt to a model serving system for execution by the machine-learned language model;
receiving, from the model serving system, an output generated by executing the machine-learned language model on the prompt, the output comprising the requested one or more SQL statements for performing the database query;
performing an SQL query included in the SQL statements;
receiving a result from performing the SQL query; and
presenting the received result as a response to the user query.
2 . The method of claim 1 , further comprising:
finetuning the machine-learned language model by using a training dataset including pairs of metadata of a data table, corresponding SQL statements that are identified from SQL dashboards, and user-submitted SQL statements that are performed on the data table.
3 . The method of claim 1 , wherein identifying an embedding for the user query to represent the user query in a latent space comprises:
generating a user query embedding based on the user query;
accessing a plurality of dashboard embeddings in a database, each dashboarding embedding representing an SQL dashboard, and each SQL dashboard comprising at least one SQL query that powers the SQL dashboard, one or more data tables corresponding to the at least one SQL query, and metadata of the SQL dashboard; and
identifying a measure of similarity between the user query embedding and at least one dashboard embedding.
4 . The method of claim 3 , wherein accessing a plurality of dashboard embeddings in a database comprise:
generating a dashboard embedding for each of a plurality of SQL dashboards using the corresponding metadata of the SQL dashboard.
5 . The method of claim 3 , wherein identifying, based on the determined measure of similarity, an SQL dashboard corresponding to the user query comprises:
selecting an SQL dashboard corresponding to a dashboard embedding that has at least a threshold measure of similarity to the SQL dashboard corresponding to the user query.
6 . The method of claim 3 , wherein identifying, based on the determined measure of similarity, an SQL dashboard corresponding to the user query comprises:
ranking the at least one dashboard embedding based on the determined measure of similarity; and
selecting the SQL dashboard based on the ranking of the corresponding dashboard embeddings.
7 . The method of claim 3 , wherein the measure of similarity between the user query embedding and at least one dashboard embedding is a distance between the user query embedding and the at least one dashboard embedding in the latent space.
8 . A computer program product comprising a non-transitory computer readable storage medium having instructions encoded thereon that, when executed by a processor, cause the processor to perform steps comprising:
receiving, from a client device, a user query from a user, the user query comprising a task for performing a database query;
identifying an embedding for the user query to represent the user query in a latent space;
generating a prompt for input to a machine-learned language model, the prompt specifying the user query and a request to generate one or more SQL statements for performing the database query on a data table, wherein the machine-learned language model is trained by:
identifying one or more SQL dashboards implemented for the online system, wherein a SQL dashboard is powered by at least one SQL query,
for each identified SQL dashboard, obtaining a training example including a pair of inputs including metadata describing the SQL dashboard, and outputs including the at least one SQL query powering the identified SQL dashboard, the metadata including one or a combination of column features of data tables associated with the SQL dashboard, headings, descriptions, user comments, or text input for the SQL dashboard,
obtaining a loss indicating a difference between the outputs of the training example and outputs obtained by applying the machine-learned language model to the inputs of the training example, and
fine-tuning parameters of the machine-learned model to reduce the loss based on the metadata obtained for the one or more SQL dashboards and the at least one SQL query powering the one or more SQL dashboards;
providing the prompt to a model serving system for execution by the machine-learned language model;
receiving, from the model serving system, an output generated by executing the machine-learned language model on the prompt, the output comprising the requested one or more SQL statements for performing the database query;
performing an SQL query included in the SQL statements;
receiving a result from performing the SQL query; and
presenting the received result as a response to the user query.
9 . The computer program product of claim 8 , wherein the instructions further causes the processor to perform steps comprising:
finetuning the machine-learned language model by using a training dataset including pairs of metadata of a data table, corresponding SQL statements that are identified from SQL dashboards, and user-submitted SQL statements that are performed on the data table.
10 . The computer program product of claim 8 , wherein the instructions to identify an embedding for the user query to represent the user query in a latent space further cause the processor to perform steps comprising:
generating a user query embedding based on the user query;
accessing a plurality of dashboard embeddings in a database, each dashboarding embedding representing an SQL dashboard, and each SQL dashboard comprising at least one SQL query that powers the SQL dashboard, one or more data tables corresponding to the at least one SQL query, and metadata of the SQL dashboard; and
identifying a measure of similarity between the user query embedding and at least one dashboard embedding.
11 . The computer program product of claim 10 , wherein the instructions to access a plurality of dashboard embeddings in a database further cause the processor to perform steps comprising:
generating a dashboard embedding for each of a plurality of SQL dashboards using the corresponding metadata of the SQL dashboard.
12 . The computer program product of claim 10 , wherein the instructions to identify, based on the determined measure of similarity, an SQL dashboard corresponding to the user query further cause the processor to perform steps comprising:
selecting an SQL dashboard corresponding to a dashboard embedding that has at least a threshold measure of similarity the SQL dashboard corresponding to the user query.
13 . The computer program product of claim 10 , wherein the instructions to identify, based on the determined measure of similarity, an SQL dashboard corresponding to the user query further cause the processor to perform steps comprising:
ranking the at least one dashboard embedding based on the determined measure of similarity; and
selecting the SQL dashboard based on the ranking of the corresponding dashboard embeddings.
14 . The computer program product of claim 10 , wherein the measure of similarity between the user query embedding and at least one dashboard embedding is a distance between the user query embedding and the at least one dashboard embedding in the latent space.
15 . A computer system comprising:
a processor; and
a non-transitory computer-readable storage medium having instructions that, when executed by the processor, cause the computer system to perform steps comprising:
receiving, from a client device, a user query from a user, the user query comprising a task for performing a database query;
identifying an embedding for the user query to represent the user query in a latent space;
generating a prompt for input to a machine-learned language model, the prompt specifying the user query and a request to generate one or more SQL statements for performing the database query on a data table, wherein the machine-learned language model is trained by:
identifying one or more SQL dashboards implemented for the online system, wherein a SQL dashboard is powered by at least one SQL query,
for each identified SQL dashboard, obtaining a training example including a pair of inputs including metadata describing the SQL dashboard, and outputs including the at least one SQL query powering the identified SQL dashboard, the metadata including one or a combination of column features of data tables associated with the SQL dashboard, headings, descriptions, user comments, or text input for the SQL dashboard,
obtaining a loss indicating a difference between the outputs of the training example and outputs obtained by applying the machine-learned language model to the inputs of the training example, and
fine-tuning parameters of the machine-learned model to reduce the loss based on the metadata obtained for the one or more SQL dashboards and the at least one SQL query powering the one or more SQL dashboards;
providing the prompt to a model serving system for execution by the machine-learned language model;
receiving, from the model serving system, an output generated by executing the machine-learned language model on the prompt, the output comprising the requested one or more SQL statements for performing the database query;
performing an SQL query included in the SQL statements;
receiving a result from performing the SQL query; and
presenting the received result as a response to the user query.
16 . The system of claim 15 , wherein the instructions further cause the processors to perform steps comprising:
finetuning the machine-learned language model by using a training dataset including pairs of metadata of a data table, corresponding SQL statements that are identified from SQL dashboards, and user-submitted SQL statements that are performed on the data table.
17 . The system of claim 15 , wherein the instructions to identify an embedding for the user query to represent the user query in a latent space further cause the processor to perform steps comprising:
generating a user query embedding based on the user query;
accessing a plurality of dashboard embeddings in a database, each dashboarding embedding representing an SQL dashboard, and each SQL dashboard comprising at least one SQL query that powers the SQL dashboard, one or more data tables corresponding to the at least one SQL query, and metadata of the SQL dashboard; and
identifying a measure of similarity between the user query embedding and at least one dashboard embedding.
18 . The system of claim 17 , wherein the instructions to access a plurality of dashboard embeddings in a database further cause the processor to perform steps comprising:
generating a dashboard embedding for each of a plurality of SQL dashboards using the corresponding metadata of the SQL dashboard.
19 . The system of claim 17 , wherein the instructions to identify, based on the determined measure of similarity, an SQL dashboard corresponding to the user query further cause the processor to perform steps comprising:
selecting an SQL dashboard corresponding to a dashboard embedding that has at least a threshold measure of similarity the SQL dashboard corresponding to the user query.
20 . The system of claim 17 , wherein the instructions to identify, based on the determined measure of similarity, an SQL dashboard corresponding to the user query further cause the processor to perform steps comprising:
ranking the at least one dashboard embedding based on the determined measure of similarity; and
selecting the SQL dashboard based on the ranking of the corresponding dashboard embeddings.