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.