@@ -48,6 +48,7 @@ class Model(Asset):
4848 function (Text, optional): model AI function. Defaults to None.
4949 url (str): URL to run the model.
5050 backend_url (str): URL of the backend.
51+ pricing (Dict, optional): model price. Defaults to None.
5152 **additional_info: Any additional Model info to be saved
5253 """
5354
@@ -61,6 +62,7 @@ def __init__(
6162 version : Optional [Text ] = None ,
6263 function : Optional [Text ] = None ,
6364 is_subscribed : bool = False ,
65+ cost : Optional [Dict ] = None ,
6466 ** additional_info ,
6567 ) -> None :
6668 """Model Init
@@ -74,9 +76,10 @@ def __init__(
7476 version (Text, optional): version of the model. Defaults to "1.0".
7577 function (Text, optional): model AI function. Defaults to None.
7678 is_subscribed (bool, optional): Is the user subscribed. Defaults to False.
79+ cost (Dict, optional): model price. Defaults to None.
7780 **additional_info: Any additional Model info to be saved
7881 """
79- super ().__init__ (id , name , description , supplier , version )
82+ super ().__init__ (id , name , description , supplier , version , cost = cost )
8083 self .api_key = api_key
8184 self .additional_info = additional_info
8285 self .url = config .MODELS_RUN_URL
@@ -264,6 +267,7 @@ def check_finetune_status(self, after_epoch: Optional[int] = None):
264267 """
265268 from aixplain .enums .asset_status import AssetStatus
266269 from aixplain .modules .finetune .status import FinetuneStatus
270+
267271 headers = {"x-api-key" : self .api_key , "Content-Type" : "application/json" }
268272 resp = None
269273 try :
@@ -274,15 +278,15 @@ def check_finetune_status(self, after_epoch: Optional[int] = None):
274278 finetune_status = AssetStatus (resp ["finetuneStatus" ])
275279 model_status = AssetStatus (resp ["modelStatus" ])
276280 logs = sorted (resp ["logs" ], key = lambda x : float (x ["epoch" ]))
277-
281+
278282 target_epoch = None
279283 if after_epoch is not None :
280284 logs = [log for log in logs if float (log ["epoch" ]) > after_epoch ]
281285 if len (logs ) > 0 :
282286 target_epoch = float (logs [0 ]["epoch" ])
283287 elif len (logs ) > 0 :
284288 target_epoch = float (logs [- 1 ]["epoch" ])
285-
289+
286290 if target_epoch is not None :
287291 log = None
288292 for log_ in logs :
@@ -294,7 +298,7 @@ def check_finetune_status(self, after_epoch: Optional[int] = None):
294298 log ["trainLoss" ] = log_ ["trainLoss" ]
295299 if log_ ["evalLoss" ] is not None :
296300 log ["evalLoss" ] = log_ ["evalLoss" ]
297-
301+
298302 status = FinetuneStatus (
299303 status = finetune_status ,
300304 model_status = model_status ,
0 commit comments