kiln_ai.adapters.fine_tune.openai_finetune
1import time 2 3import openai 4from openai.types.fine_tuning import FineTuningJob 5 6from kiln_ai.adapters.fine_tune.base_finetune import ( 7 BaseFinetuneAdapter, 8 FineTuneParameter, 9 FineTuneStatus, 10 FineTuneStatusType, 11) 12from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetFormatter 13from kiln_ai.datamodel import DatasetSplit, Task 14from kiln_ai.utils.config import Config 15 16oai_client = openai.OpenAI( 17 api_key=Config.shared().open_ai_api_key or "", 18) 19 20 21class OpenAIFinetune(BaseFinetuneAdapter): 22 """ 23 A fine-tuning adapter for OpenAI. 24 """ 25 26 def status(self) -> FineTuneStatus: 27 if not self.datamodel or not self.datamodel.provider_id: 28 return FineTuneStatus( 29 status=FineTuneStatusType.pending, 30 message="This fine-tune has not been started or has not been assigned a provider ID.", 31 ) 32 33 try: 34 # Will raise an error if the job is not found, or for other issues 35 response = oai_client.fine_tuning.jobs.retrieve(self.datamodel.provider_id) 36 except openai.APIConnectionError: 37 return FineTuneStatus( 38 status=FineTuneStatusType.unknown, message="Server connection error" 39 ) 40 except openai.RateLimitError: 41 return FineTuneStatus( 42 status=FineTuneStatusType.unknown, 43 message="Rate limit exceeded. Could not fetch fine-tune status.", 44 ) 45 except openai.APIStatusError as e: 46 if e.status_code == 404: 47 return FineTuneStatus( 48 status=FineTuneStatusType.unknown, 49 message="Job with this ID not found. It may have been deleted.", 50 ) 51 return FineTuneStatus( 52 status=FineTuneStatusType.unknown, 53 message=f"Unknown error: [{str(e)}]", 54 ) 55 56 if not response or not isinstance(response, FineTuningJob): 57 return FineTuneStatus( 58 status=FineTuneStatusType.unknown, 59 message="Invalid response from OpenAI", 60 ) 61 if response.error and response.error.code: 62 return FineTuneStatus( 63 status=FineTuneStatusType.failed, 64 message=f"{response.error.message} [Code: {response.error.code}]", 65 ) 66 status = response.status 67 if status == "failed": 68 return FineTuneStatus( 69 status=FineTuneStatusType.failed, 70 message="Job failed - unknown reason", 71 ) 72 if status == "cancelled": 73 return FineTuneStatus( 74 status=FineTuneStatusType.failed, message="Job cancelled" 75 ) 76 if status in ["validating_files", "running", "queued"]: 77 time_to_finish_msg: str | None = None 78 if response.estimated_finish is not None: 79 time_to_finish_msg = f"Estimated finish time: {int(response.estimated_finish - time.time())} seconds." 80 return FineTuneStatus( 81 status=FineTuneStatusType.running, 82 message=f"Job is still running [{status}]. {time_to_finish_msg or ''}", 83 ) 84 if status == "succeeded": 85 return FineTuneStatus( 86 status=FineTuneStatusType.completed, message="Training job completed" 87 ) 88 return FineTuneStatus( 89 status=FineTuneStatusType.unknown, 90 message=f"Unknown status: [{status}]", 91 ) 92 93 def _start(self, dataset: DatasetSplit) -> None: 94 task = self.datamodel.parent_task() 95 if not task: 96 raise ValueError("Task is required to start a fine-tune") 97 98 train_file_id = self.generate_and_upload_jsonl( 99 dataset, self.datamodel.train_split_name, task 100 ) 101 validation_file_id = None 102 if self.datamodel.validation_split_name: 103 validation_file_id = self.generate_and_upload_jsonl( 104 dataset, self.datamodel.validation_split_name, task 105 ) 106 107 hyperparameters = { 108 k: v 109 for k, v in self.datamodel.parameters.items() 110 if k in ["n_epochs", "learning_rate_multiplier", "batch_size"] 111 } 112 113 ft = oai_client.fine_tuning.jobs.create( 114 training_file=train_file_id, 115 model=self.datamodel.base_model_id, 116 validation_file=validation_file_id, 117 seed=self.datamodel.parameters.get("seed"), # type: ignore 118 hyperparameters=hyperparameters, # type: ignore 119 suffix=f"kiln_ai.{self.datamodel.id}", 120 ) 121 self.datamodel.provider_id = ft.id 122 # Model can get more specific after fine-tune call (gpt-4o-mini to gpt-4o-mini-2024-07-18) 123 self.datamodel.base_model_id = ft.model 124 125 return None 126 127 def generate_and_upload_jsonl( 128 self, dataset: DatasetSplit, split_name: str, task: Task 129 ) -> str: 130 formatter = DatasetFormatter(dataset, self.datamodel.system_message) 131 # All OpenAI models support tool calls for structured outputs 132 format = ( 133 DatasetFormat.CHAT_MESSAGE_TOOLCALL_JSONL 134 if task.output_json_schema 135 else DatasetFormat.CHAT_MESSAGE_RESPONSE_JSONL 136 ) 137 path = formatter.dump_to_file(split_name, format) 138 139 response = oai_client.files.create( 140 file=open(path, "rb"), 141 purpose="fine-tune", 142 ) 143 id = response.id 144 if not id: 145 raise ValueError("Failed to upload file to OpenAI") 146 return id 147 148 @classmethod 149 def available_parameters(cls) -> list[FineTuneParameter]: 150 return [ 151 FineTuneParameter( 152 name="batch_size", 153 type="int", 154 description="Number of examples in each batch. A larger batch size means that model parameters are updated less frequently, but with lower variance. Defaults to 'auto'", 155 ), 156 FineTuneParameter( 157 name="learning_rate_multiplier", 158 type="float", 159 description="Scaling factor for the learning rate. A smaller learning rate may be useful to avoid overfitting. Defaults to 'auto'", 160 optional=True, 161 ), 162 FineTuneParameter( 163 name="n_epochs", 164 type="int", 165 description="The number of epochs to train the model for. An epoch refers to one full cycle through the training dataset. Defaults to 'auto'", 166 optional=True, 167 ), 168 FineTuneParameter( 169 name="seed", 170 type="int", 171 description="The seed controls the reproducibility of the job. Passing in the same seed and job parameters should produce the same results, but may differ in rare cases. If a seed is not specified, one will be generated for you.", 172 optional=True, 173 ), 174 ]
oai_client =
<openai.OpenAI object>
22class OpenAIFinetune(BaseFinetuneAdapter): 23 """ 24 A fine-tuning adapter for OpenAI. 25 """ 26 27 def status(self) -> FineTuneStatus: 28 if not self.datamodel or not self.datamodel.provider_id: 29 return FineTuneStatus( 30 status=FineTuneStatusType.pending, 31 message="This fine-tune has not been started or has not been assigned a provider ID.", 32 ) 33 34 try: 35 # Will raise an error if the job is not found, or for other issues 36 response = oai_client.fine_tuning.jobs.retrieve(self.datamodel.provider_id) 37 except openai.APIConnectionError: 38 return FineTuneStatus( 39 status=FineTuneStatusType.unknown, message="Server connection error" 40 ) 41 except openai.RateLimitError: 42 return FineTuneStatus( 43 status=FineTuneStatusType.unknown, 44 message="Rate limit exceeded. Could not fetch fine-tune status.", 45 ) 46 except openai.APIStatusError as e: 47 if e.status_code == 404: 48 return FineTuneStatus( 49 status=FineTuneStatusType.unknown, 50 message="Job with this ID not found. It may have been deleted.", 51 ) 52 return FineTuneStatus( 53 status=FineTuneStatusType.unknown, 54 message=f"Unknown error: [{str(e)}]", 55 ) 56 57 if not response or not isinstance(response, FineTuningJob): 58 return FineTuneStatus( 59 status=FineTuneStatusType.unknown, 60 message="Invalid response from OpenAI", 61 ) 62 if response.error and response.error.code: 63 return FineTuneStatus( 64 status=FineTuneStatusType.failed, 65 message=f"{response.error.message} [Code: {response.error.code}]", 66 ) 67 status = response.status 68 if status == "failed": 69 return FineTuneStatus( 70 status=FineTuneStatusType.failed, 71 message="Job failed - unknown reason", 72 ) 73 if status == "cancelled": 74 return FineTuneStatus( 75 status=FineTuneStatusType.failed, message="Job cancelled" 76 ) 77 if status in ["validating_files", "running", "queued"]: 78 time_to_finish_msg: str | None = None 79 if response.estimated_finish is not None: 80 time_to_finish_msg = f"Estimated finish time: {int(response.estimated_finish - time.time())} seconds." 81 return FineTuneStatus( 82 status=FineTuneStatusType.running, 83 message=f"Job is still running [{status}]. {time_to_finish_msg or ''}", 84 ) 85 if status == "succeeded": 86 return FineTuneStatus( 87 status=FineTuneStatusType.completed, message="Training job completed" 88 ) 89 return FineTuneStatus( 90 status=FineTuneStatusType.unknown, 91 message=f"Unknown status: [{status}]", 92 ) 93 94 def _start(self, dataset: DatasetSplit) -> None: 95 task = self.datamodel.parent_task() 96 if not task: 97 raise ValueError("Task is required to start a fine-tune") 98 99 train_file_id = self.generate_and_upload_jsonl( 100 dataset, self.datamodel.train_split_name, task 101 ) 102 validation_file_id = None 103 if self.datamodel.validation_split_name: 104 validation_file_id = self.generate_and_upload_jsonl( 105 dataset, self.datamodel.validation_split_name, task 106 ) 107 108 hyperparameters = { 109 k: v 110 for k, v in self.datamodel.parameters.items() 111 if k in ["n_epochs", "learning_rate_multiplier", "batch_size"] 112 } 113 114 ft = oai_client.fine_tuning.jobs.create( 115 training_file=train_file_id, 116 model=self.datamodel.base_model_id, 117 validation_file=validation_file_id, 118 seed=self.datamodel.parameters.get("seed"), # type: ignore 119 hyperparameters=hyperparameters, # type: ignore 120 suffix=f"kiln_ai.{self.datamodel.id}", 121 ) 122 self.datamodel.provider_id = ft.id 123 # Model can get more specific after fine-tune call (gpt-4o-mini to gpt-4o-mini-2024-07-18) 124 self.datamodel.base_model_id = ft.model 125 126 return None 127 128 def generate_and_upload_jsonl( 129 self, dataset: DatasetSplit, split_name: str, task: Task 130 ) -> str: 131 formatter = DatasetFormatter(dataset, self.datamodel.system_message) 132 # All OpenAI models support tool calls for structured outputs 133 format = ( 134 DatasetFormat.CHAT_MESSAGE_TOOLCALL_JSONL 135 if task.output_json_schema 136 else DatasetFormat.CHAT_MESSAGE_RESPONSE_JSONL 137 ) 138 path = formatter.dump_to_file(split_name, format) 139 140 response = oai_client.files.create( 141 file=open(path, "rb"), 142 purpose="fine-tune", 143 ) 144 id = response.id 145 if not id: 146 raise ValueError("Failed to upload file to OpenAI") 147 return id 148 149 @classmethod 150 def available_parameters(cls) -> list[FineTuneParameter]: 151 return [ 152 FineTuneParameter( 153 name="batch_size", 154 type="int", 155 description="Number of examples in each batch. A larger batch size means that model parameters are updated less frequently, but with lower variance. Defaults to 'auto'", 156 ), 157 FineTuneParameter( 158 name="learning_rate_multiplier", 159 type="float", 160 description="Scaling factor for the learning rate. A smaller learning rate may be useful to avoid overfitting. Defaults to 'auto'", 161 optional=True, 162 ), 163 FineTuneParameter( 164 name="n_epochs", 165 type="int", 166 description="The number of epochs to train the model for. An epoch refers to one full cycle through the training dataset. Defaults to 'auto'", 167 optional=True, 168 ), 169 FineTuneParameter( 170 name="seed", 171 type="int", 172 description="The seed controls the reproducibility of the job. Passing in the same seed and job parameters should produce the same results, but may differ in rare cases. If a seed is not specified, one will be generated for you.", 173 optional=True, 174 ), 175 ]
A fine-tuning adapter for OpenAI.
27 def status(self) -> FineTuneStatus: 28 if not self.datamodel or not self.datamodel.provider_id: 29 return FineTuneStatus( 30 status=FineTuneStatusType.pending, 31 message="This fine-tune has not been started or has not been assigned a provider ID.", 32 ) 33 34 try: 35 # Will raise an error if the job is not found, or for other issues 36 response = oai_client.fine_tuning.jobs.retrieve(self.datamodel.provider_id) 37 except openai.APIConnectionError: 38 return FineTuneStatus( 39 status=FineTuneStatusType.unknown, message="Server connection error" 40 ) 41 except openai.RateLimitError: 42 return FineTuneStatus( 43 status=FineTuneStatusType.unknown, 44 message="Rate limit exceeded. Could not fetch fine-tune status.", 45 ) 46 except openai.APIStatusError as e: 47 if e.status_code == 404: 48 return FineTuneStatus( 49 status=FineTuneStatusType.unknown, 50 message="Job with this ID not found. It may have been deleted.", 51 ) 52 return FineTuneStatus( 53 status=FineTuneStatusType.unknown, 54 message=f"Unknown error: [{str(e)}]", 55 ) 56 57 if not response or not isinstance(response, FineTuningJob): 58 return FineTuneStatus( 59 status=FineTuneStatusType.unknown, 60 message="Invalid response from OpenAI", 61 ) 62 if response.error and response.error.code: 63 return FineTuneStatus( 64 status=FineTuneStatusType.failed, 65 message=f"{response.error.message} [Code: {response.error.code}]", 66 ) 67 status = response.status 68 if status == "failed": 69 return FineTuneStatus( 70 status=FineTuneStatusType.failed, 71 message="Job failed - unknown reason", 72 ) 73 if status == "cancelled": 74 return FineTuneStatus( 75 status=FineTuneStatusType.failed, message="Job cancelled" 76 ) 77 if status in ["validating_files", "running", "queued"]: 78 time_to_finish_msg: str | None = None 79 if response.estimated_finish is not None: 80 time_to_finish_msg = f"Estimated finish time: {int(response.estimated_finish - time.time())} seconds." 81 return FineTuneStatus( 82 status=FineTuneStatusType.running, 83 message=f"Job is still running [{status}]. {time_to_finish_msg or ''}", 84 ) 85 if status == "succeeded": 86 return FineTuneStatus( 87 status=FineTuneStatusType.completed, message="Training job completed" 88 ) 89 return FineTuneStatus( 90 status=FineTuneStatusType.unknown, 91 message=f"Unknown status: [{status}]", 92 )
Get the status of the fine-tune.
def
generate_and_upload_jsonl( self, dataset: kiln_ai.datamodel.DatasetSplit, split_name: str, task: kiln_ai.datamodel.Task) -> str:
128 def generate_and_upload_jsonl( 129 self, dataset: DatasetSplit, split_name: str, task: Task 130 ) -> str: 131 formatter = DatasetFormatter(dataset, self.datamodel.system_message) 132 # All OpenAI models support tool calls for structured outputs 133 format = ( 134 DatasetFormat.CHAT_MESSAGE_TOOLCALL_JSONL 135 if task.output_json_schema 136 else DatasetFormat.CHAT_MESSAGE_RESPONSE_JSONL 137 ) 138 path = formatter.dump_to_file(split_name, format) 139 140 response = oai_client.files.create( 141 file=open(path, "rb"), 142 purpose="fine-tune", 143 ) 144 id = response.id 145 if not id: 146 raise ValueError("Failed to upload file to OpenAI") 147 return id
@classmethod
def
available_parameters(cls) -> list[kiln_ai.adapters.fine_tune.base_finetune.FineTuneParameter]:
149 @classmethod 150 def available_parameters(cls) -> list[FineTuneParameter]: 151 return [ 152 FineTuneParameter( 153 name="batch_size", 154 type="int", 155 description="Number of examples in each batch. A larger batch size means that model parameters are updated less frequently, but with lower variance. Defaults to 'auto'", 156 ), 157 FineTuneParameter( 158 name="learning_rate_multiplier", 159 type="float", 160 description="Scaling factor for the learning rate. A smaller learning rate may be useful to avoid overfitting. Defaults to 'auto'", 161 optional=True, 162 ), 163 FineTuneParameter( 164 name="n_epochs", 165 type="int", 166 description="The number of epochs to train the model for. An epoch refers to one full cycle through the training dataset. Defaults to 'auto'", 167 optional=True, 168 ), 169 FineTuneParameter( 170 name="seed", 171 type="int", 172 description="The seed controls the reproducibility of the job. Passing in the same seed and job parameters should produce the same results, but may differ in rare cases. If a seed is not specified, one will be generated for you.", 173 optional=True, 174 ), 175 ]
Returns a list of parameters that can be provided for this fine-tune.