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, StructuredOutputMode, Task
 14from kiln_ai.utils.config import Config
 15
 16oai_client = openai.AsyncOpenAI(
 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    async def status(self) -> FineTuneStatus:
 27        """
 28        Get the status of the fine-tune.
 29        """
 30
 31        # Update the datamodel with the latest status if it has changed
 32        status = await self._status()
 33        if status.status != self.datamodel.latest_status:
 34            self.datamodel.latest_status = status.status
 35            if self.datamodel.path:
 36                self.datamodel.save_to_file()
 37        return status
 38
 39    async def _status(self) -> FineTuneStatus:
 40        if not self.datamodel or not self.datamodel.provider_id:
 41            return FineTuneStatus(
 42                status=FineTuneStatusType.pending,
 43                message="This fine-tune has not been started or has not been assigned a provider ID.",
 44            )
 45
 46        try:
 47            # Will raise an error if the job is not found, or for other issues
 48            response = await oai_client.fine_tuning.jobs.retrieve(
 49                self.datamodel.provider_id
 50            )
 51
 52            # If the fine-tuned model has been updated, update the datamodel
 53            try:
 54                if (
 55                    self.datamodel.fine_tune_model_id != response.fine_tuned_model
 56                    or self.datamodel.base_model_id != response.model
 57                ):
 58                    self.datamodel.fine_tune_model_id = response.fine_tuned_model
 59                    self.datamodel.base_model_id = response.model
 60                    self.datamodel.save_to_file()
 61            except Exception:
 62                # Don't let this error crash the status call
 63                pass
 64
 65        except openai.APIConnectionError:
 66            return FineTuneStatus(
 67                status=FineTuneStatusType.unknown, message="Server connection error"
 68            )
 69        except openai.RateLimitError:
 70            return FineTuneStatus(
 71                status=FineTuneStatusType.unknown,
 72                message="Rate limit exceeded. Could not fetch fine-tune status.",
 73            )
 74        except openai.APIStatusError as e:
 75            if e.status_code == 404:
 76                return FineTuneStatus(
 77                    status=FineTuneStatusType.unknown,
 78                    message="Job with this ID not found. It may have been deleted.",
 79                )
 80            return FineTuneStatus(
 81                status=FineTuneStatusType.unknown,
 82                message=f"Unknown error: [{str(e)}]",
 83            )
 84
 85        if not response or not isinstance(response, FineTuningJob):
 86            return FineTuneStatus(
 87                status=FineTuneStatusType.unknown,
 88                message="Invalid response from OpenAI",
 89            )
 90        if response.error and response.error.code:
 91            return FineTuneStatus(
 92                status=FineTuneStatusType.failed,
 93                message=f"{response.error.message} [Code: {response.error.code}]",
 94            )
 95        status = response.status
 96        if status == "failed":
 97            return FineTuneStatus(
 98                status=FineTuneStatusType.failed,
 99                message="Job failed - unknown reason",
100            )
101        if status == "cancelled":
102            return FineTuneStatus(
103                status=FineTuneStatusType.failed, message="Job cancelled"
104            )
105        if status in ["validating_files", "running", "queued"]:
106            time_to_finish_msg: str | None = None
107            if response.estimated_finish is not None:
108                time_to_finish_msg = f"Estimated finish time: {int(response.estimated_finish - time.time())} seconds."
109            return FineTuneStatus(
110                status=FineTuneStatusType.running,
111                message=f"Fine tune job is running [{status}]. {time_to_finish_msg or ''}",
112            )
113        if status == "succeeded":
114            return FineTuneStatus(
115                status=FineTuneStatusType.completed, message="Training job completed"
116            )
117        return FineTuneStatus(
118            status=FineTuneStatusType.unknown,
119            message=f"Unknown status: [{status}]",
120        )
121
122    async def _start(self, dataset: DatasetSplit) -> None:
123        task = self.datamodel.parent_task()
124        if not task:
125            raise ValueError("Task is required to start a fine-tune")
126
127        # Use chat format for unstructured output, and JSON for formatted output (was previously function calls)
128        format = DatasetFormat.OPENAI_CHAT_JSONL
129        if task.output_json_schema:
130            format = DatasetFormat.OPENAI_CHAT_JSON_SCHEMA_JSONL
131            self.datamodel.structured_output_mode = StructuredOutputMode.json_schema
132        train_file_id = await self.generate_and_upload_jsonl(
133            dataset, self.datamodel.train_split_name, task, format
134        )
135        validation_file_id = None
136        if self.datamodel.validation_split_name:
137            validation_file_id = await self.generate_and_upload_jsonl(
138                dataset, self.datamodel.validation_split_name, task, format
139            )
140
141        # Filter to hyperparameters which are set via the hyperparameters field (some like seed are set via the API)
142        hyperparameters = {
143            k: v
144            for k, v in self.datamodel.parameters.items()
145            if k in ["n_epochs", "learning_rate_multiplier", "batch_size"]
146        }
147
148        ft = await oai_client.fine_tuning.jobs.create(
149            training_file=train_file_id,
150            model=self.datamodel.base_model_id,
151            validation_file=validation_file_id,
152            seed=self.datamodel.parameters.get("seed"),  # type: ignore
153            hyperparameters=hyperparameters,  # type: ignore
154            suffix=f"kiln_ai.{self.datamodel.id}",
155        )
156        self.datamodel.provider_id = ft.id
157        self.datamodel.fine_tune_model_id = ft.fine_tuned_model
158        # Model can get more specific after fine-tune call (gpt-4o-mini to gpt-4o-mini-2024-07-18) so we update it in the datamodel
159        self.datamodel.base_model_id = ft.model
160
161        return None
162
163    async def generate_and_upload_jsonl(
164        self, dataset: DatasetSplit, split_name: str, task: Task, format: DatasetFormat
165    ) -> str:
166        formatter = DatasetFormatter(
167            dataset, self.datamodel.system_message, self.datamodel.thinking_instructions
168        )
169        path = formatter.dump_to_file(split_name, format, self.datamodel.data_strategy)
170
171        response = await oai_client.files.create(
172            file=open(path, "rb"),
173            purpose="fine-tune",
174        )
175        id = response.id
176        if not id:
177            raise ValueError("Failed to upload file to OpenAI")
178        return id
179
180    @classmethod
181    def available_parameters(cls) -> list[FineTuneParameter]:
182        return [
183            FineTuneParameter(
184                name="batch_size",
185                type="int",
186                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'",
187            ),
188            FineTuneParameter(
189                name="learning_rate_multiplier",
190                type="float",
191                description="Scaling factor for the learning rate. A smaller learning rate may be useful to avoid overfitting. Defaults to 'auto'",
192                optional=True,
193            ),
194            FineTuneParameter(
195                name="n_epochs",
196                type="int",
197                description="The number of epochs to train the model for. An epoch refers to one full cycle through the training dataset. Defaults to 'auto'",
198                optional=True,
199            ),
200            FineTuneParameter(
201                name="seed",
202                type="int",
203                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.",
204                optional=True,
205            ),
206        ]