@@ -42,7 +42,7 @@ class Status(Enum):
4242 FAILED = "FAILED"
4343
4444 def upsert_labels (self , label_ids , timeout_seconds = 60 ):
45- """ Adds data rows and labels to a model run
45+ """ Adds data rows and labels to a Model Run
4646 Args:
4747 label_ids (list): label ids to insert
4848 timeout_seconds (float): Max waiting time, in seconds.
@@ -75,7 +75,7 @@ def upsert_labels(self, label_ids, timeout_seconds=60):
7575 timeout_seconds = timeout_seconds )
7676
7777 def upsert_data_rows (self , data_row_ids , timeout_seconds = 60 ):
78- """ Adds data rows to a model run without any associated labels
78+ """ Adds data rows to a Model Run without any associated labels
7979 Args:
8080 data_row_ids (list): data row ids to add to mea
8181 timeout_seconds (float): Max waiting time, in seconds.
@@ -167,7 +167,7 @@ def model_run_data_rows(self):
167167 ['annotationGroups' , 'pageInfo' , 'endCursor' ])
168168
169169 def delete (self ):
170- """ Deletes specified model run .
170+ """ Deletes specified Model Run .
171171
172172 Returns:
173173 Query execution success.
@@ -178,10 +178,10 @@ def delete(self):
178178 self .client .execute (query_str , {ids_param : str (self .uid )})
179179
180180 def delete_model_run_data_rows (self , data_row_ids : List [str ]):
181- """ Deletes data rows from model runs .
181+ """ Deletes data rows from Model Runs .
182182
183183 Args:
184- data_row_ids (list): List of data row ids to delete from the model run .
184+ data_row_ids (list): List of data row ids to delete from the Model Run .
185185 Returns:
186186 Query execution success.
187187 """
@@ -262,6 +262,56 @@ def update_status(self,
262262 },
263263 experimental = True )
264264
265+ @experimental
266+ def update_config (self , config : Dict [str , Any ]) -> Dict [str , Any ]:
267+ """
268+ Updates the Model Run's training metadata config
269+ Args:
270+ config (dict): A dictionary of keys and values
271+ Returns:
272+ Model Run id and updated training metadata
273+ """
274+ data : Dict [str , Any ] = {'config' : config }
275+ res = self .client .execute (
276+ """mutation updateModelRunConfigPyApi($modelRunId: ID!, $data: UpdateModelRunConfigInput!){
277+ updateModelRunConfig(modelRun: {id : $modelRunId}, data: $data){trainingMetadata}
278+ }
279+ """ , {
280+ 'modelRunId' : self .uid ,
281+ 'data' : data
282+ },
283+ experimental = True )
284+ return res ["updateModelRunConfig" ]
285+
286+ @experimental
287+ def reset_config (self ) -> Dict [str , Any ]:
288+ """
289+ Resets Model Run's training metadata config
290+ Returns:
291+ Model Run id and reset training metadata
292+ """
293+ res = self .client .execute (
294+ """mutation resetModelRunConfigPyApi($modelRunId: ID!){
295+ resetModelRunConfig(modelRun: {id : $modelRunId}){trainingMetadata}
296+ }
297+ """ , {'modelRunId' : self .uid },
298+ experimental = True )
299+ return res ["resetModelRunConfig" ]
300+
301+ @experimental
302+ def get_config (self ) -> Dict [str , Any ]:
303+ """
304+ Gets Model Run's training metadata
305+ Returns:
306+ training metadata as a dictionary
307+ """
308+ res = self .client .execute ("""query ModelRunPyApi($modelRunId: ID!){
309+ modelRun(where: {id : $modelRunId}){trainingMetadata}
310+ }
311+ """ , {'modelRunId' : self .uid },
312+ experimental = True )
313+ return res ["modelRun" ]["trainingMetadata" ]
314+
265315 @experimental
266316 def export_labels (
267317 self ,
0 commit comments