@@ -32,7 +32,7 @@ def status(self):
32
32
except ImportError :
33
33
raise ImportError (
34
34
"To use Databricks finetuning, please install the databricks_genai package via "
35
- "`pip install databricks_genai`."
35
+ "`pip install databricks_genai`." ,
36
36
)
37
37
run = fm .get (self .finetuning_run )
38
38
return run .status
@@ -84,7 +84,7 @@ def deploy_finetuned_model(
84
84
model_name = model .replace ("." , "_" )
85
85
86
86
get_endpoint_response = requests .get (
87
- url = f"{ databricks_host } /api/2.0/serving-endpoints/{ model_name } " , json = {"name" : model_name }, headers = headers
87
+ url = f"{ databricks_host } /api/2.0/serving-endpoints/{ model_name } " , json = {"name" : model_name }, headers = headers ,
88
88
)
89
89
90
90
if get_endpoint_response .status_code == 200 :
@@ -98,8 +98,8 @@ def deploy_finetuned_model(
98
98
"entity_version" : model_version ,
99
99
"min_provisioned_throughput" : min_provisioned_throughput ,
100
100
"max_provisioned_throughput" : max_provisioned_throughput ,
101
- }
102
- ]
101
+ },
102
+ ],
103
103
}
104
104
105
105
response = requests .put (
@@ -120,23 +120,23 @@ def deploy_finetuned_model(
120
120
"entity_version" : model_version ,
121
121
"min_provisioned_throughput" : min_provisioned_throughput ,
122
122
"max_provisioned_throughput" : max_provisioned_throughput ,
123
- }
124
- ]
123
+ },
124
+ ],
125
125
},
126
126
}
127
127
128
128
response = requests .post (url = f"{ databricks_host } /api/2.0/serving-endpoints" , json = data , headers = headers )
129
129
130
130
if response .status_code == 200 :
131
131
logger .info (
132
- f"Successfully started creating/updating serving endpoint { model_name } on Databricks model serving!"
132
+ f"Successfully started creating/updating serving endpoint { model_name } on Databricks model serving!" ,
133
133
)
134
134
else :
135
135
raise ValueError (f"Failed to create serving endpoint: { response .json ()} ." )
136
136
137
137
logger .info (
138
138
f"Waiting for serving endpoint { model_name } to be ready, this might take a few minutes... You can check "
139
- f"the status of the endpoint at { databricks_host } /ml/endpoints/{ model_name } "
139
+ f"the status of the endpoint at { databricks_host } /ml/endpoints/{ model_name } " ,
140
140
)
141
141
from openai import OpenAI
142
142
@@ -150,7 +150,7 @@ def deploy_finetuned_model(
150
150
try :
151
151
if data_format == TrainDataFormat .CHAT :
152
152
client .chat .completions .create (
153
- messages = [{"role" : "user" , "content" : "hi" }], model = model_name , max_tokens = 1
153
+ messages = [{"role" : "user" , "content" : "hi" }], model = model_name , max_tokens = 1 ,
154
154
)
155
155
elif data_format == TrainDataFormat .COMPLETION :
156
156
client .completions .create (prompt = "hi" , model = model_name , max_tokens = 1 )
@@ -161,7 +161,7 @@ def deploy_finetuned_model(
161
161
162
162
raise ValueError (
163
163
f"Failed to create serving endpoint { model_name } on Databricks model serving platform within "
164
- f"{ deploy_timeout } seconds."
164
+ f"{ deploy_timeout } seconds." ,
165
165
)
166
166
167
167
@staticmethod
@@ -179,22 +179,22 @@ def finetune(
179
179
train_data_format = TrainDataFormat .COMPLETION
180
180
else :
181
181
raise ValueError (
182
- f"String `train_data_format` must be one of 'chat' or 'completion', but received: { train_data_format } ."
182
+ f"String `train_data_format` must be one of 'chat' or 'completion', but received: { train_data_format } ." ,
183
183
)
184
184
185
185
if "train_data_path" not in train_kwargs :
186
186
raise ValueError ("The `train_data_path` must be provided to finetune on Databricks." )
187
187
# Add the file name to the directory path.
188
188
train_kwargs ["train_data_path" ] = DatabricksProvider .upload_data (
189
- train_data , train_kwargs ["train_data_path" ], train_data_format
189
+ train_data , train_kwargs ["train_data_path" ], train_data_format ,
190
190
)
191
191
192
192
try :
193
193
from databricks .model_training import foundation_model as fm
194
194
except ImportError :
195
195
raise ImportError (
196
196
"To use Databricks finetuning, please install the databricks_genai package via "
197
- "`pip install databricks_genai`."
197
+ "`pip install databricks_genai`." ,
198
198
)
199
199
200
200
if "register_to" not in train_kwargs :
@@ -224,7 +224,7 @@ def finetune(
224
224
elif job .run .status .display_name == "Failed" :
225
225
raise ValueError (
226
226
f"Finetuning run failed with status: { job .run .status .display_name } . Please check the Databricks "
227
- f"workspace for more details. Finetuning job's metadata: { job .run } ."
227
+ f"workspace for more details. Finetuning job's metadata: { job .run } ." ,
228
228
)
229
229
else :
230
230
time .sleep (60 )
@@ -236,7 +236,7 @@ def finetune(
236
236
model_to_deploy = train_kwargs .get ("register_to" )
237
237
job .endpoint_name = model_to_deploy .replace ("." , "_" )
238
238
DatabricksProvider .deploy_finetuned_model (
239
- model_to_deploy , train_data_format , databricks_host , databricks_token , deploy_timeout
239
+ model_to_deploy , train_data_format , databricks_host , databricks_token , deploy_timeout ,
240
240
)
241
241
job .launch_completed = True
242
242
# The finetuned model name should be in the format: "databricks/<endpoint_name>".
@@ -266,7 +266,7 @@ def _get_workspace_client() -> "WorkspaceClient":
266
266
except ImportError :
267
267
raise ImportError (
268
268
"To use Databricks finetuning, please install the databricks-sdk package via "
269
- "`pip install databricks-sdk`."
269
+ "`pip install databricks-sdk`." ,
270
270
)
271
271
return WorkspaceClient ()
272
272
@@ -277,7 +277,7 @@ def _create_directory_in_databricks_unity_catalog(w: "WorkspaceClient", databric
277
277
if not match :
278
278
raise ValueError (
279
279
f"Databricks Unity Catalog path must be in the format '/Volumes/<catalog>/<schema>/<volume>/...', but "
280
- f"received: { databricks_unity_catalog_path } ."
280
+ f"received: { databricks_unity_catalog_path } ." ,
281
281
)
282
282
283
283
catalog = match .group ("catalog" )
@@ -290,7 +290,7 @@ def _create_directory_in_databricks_unity_catalog(w: "WorkspaceClient", databric
290
290
except Exception :
291
291
raise ValueError (
292
292
f"Databricks Unity Catalog volume does not exist: { volume_path } , please create it on the Databricks "
293
- "workspace."
293
+ "workspace." ,
294
294
)
295
295
296
296
try :
@@ -326,32 +326,32 @@ def _validate_chat_data(data: dict[str, Any]):
326
326
if "messages" not in data :
327
327
raise ValueError (
328
328
"Each finetuning data must be a dict with a 'messages' key when `task=CHAT_COMPLETION`, but "
329
- f"received: { data } "
329
+ f"received: { data } " ,
330
330
)
331
331
332
332
if not isinstance (data ["messages" ], list ):
333
333
raise ValueError (
334
334
"The value of the 'messages' key in each finetuning data must be a list of dicts with keys 'role' and "
335
- f"'content' when `task=CHAT_COMPLETION`, but received: { data ['messages' ]} "
335
+ f"'content' when `task=CHAT_COMPLETION`, but received: { data ['messages' ]} " ,
336
336
)
337
337
338
338
for message in data ["messages" ]:
339
339
if "role" not in message :
340
340
raise ValueError (f"Each message in the 'messages' list must contain a 'role' key, but received: { message } ." )
341
341
if "content" not in message :
342
342
raise ValueError (
343
- f"Each message in the 'messages' list must contain a 'content' key, but received: { message } ."
343
+ f"Each message in the 'messages' list must contain a 'content' key, but received: { message } ." ,
344
344
)
345
345
346
346
347
347
def _validate_completion_data (data : dict [str , Any ]):
348
348
if "prompt" not in data :
349
349
raise ValueError (
350
350
"Each finetuning data must be a dict with a 'prompt' key when `task=INSTRUCTION_FINETUNE`, but "
351
- f"received: { data } "
351
+ f"received: { data } " ,
352
352
)
353
353
if "response" not in data and "completion" not in data :
354
354
raise ValueError (
355
355
"Each finetuning data must be a dict with a 'response' or 'completion' key when "
356
- f"`task=INSTRUCTION_FINETUNE`, but received: { data } "
356
+ f"`task=INSTRUCTION_FINETUNE`, but received: { data } " ,
357
357
)
0 commit comments