Finetune fix (#7629)

* Update local fine-tuning

* Make train_data_format a required argument for finetune

* Update finetune optimizers

* Update finetune classification demo

* Update launch check

* Update lm_local.py

* Revert notebook. Revert bootstrap_trace_data.

* Add back launch_lms etc

---------

Co-authored-by: Omar Khattab <[email protected]>

Author: dilarasoylu

dspy/clients/databricks.py

@@ -8,7 +8,7 @@
 import ujson
 
 from dspy.clients.provider import Provider, TrainingJob
-from dspy.clients.utils_finetune import DataFormat, get_finetune_directory
+from dspy.clients.utils_finetune import TrainDataFormat, get_finetune_directory
 
 if TYPE_CHECKING:
     from databricks.sdk import WorkspaceClient
@@ -50,7 +50,7 @@ def is_provider_model(model: str) -> bool:
     @staticmethod
     def deploy_finetuned_model(
         model: str,
-        data_format: Optional[DataFormat] = None,
+        data_format: Optional[TrainDataFormat] = None,
         databricks_host: Optional[str] = None,
         databricks_token: Optional[str] = None,
         deploy_timeout: int = 900,
@@ -148,11 +148,11 @@ def deploy_finetuned_model(
         num_retries = deploy_timeout // 60
         for _ in range(num_retries):
             try:
-                if data_format == DataFormat.chat:
+                if data_format == TrainDataFormat.CHAT:
                     client.chat.completions.create(
                         messages=[{"role": "user", "content": "hi"}], model=model_name, max_tokens=1
                     )
-                elif data_format == DataFormat.completion:
+                elif data_format == TrainDataFormat.COMPLETION:
                     client.completions.create(prompt="hi", model=model_name, max_tokens=1)
                 logger.info(f"Databricks model serving endpoint {model_name} is ready!")
                 return
@@ -169,17 +169,17 @@ def finetune(
         job: TrainingJobDatabricks,
         model: str,
         train_data: List[Dict[str, Any]],
+        train_data_format: Optional[Union[TrainDataFormat, str]] = "chat",
         train_kwargs: Optional[Dict[str, Any]] = None,
-        data_format: Optional[Union[DataFormat, str]] = "chat",
     ) -> str:
         if isinstance(data_format, str):
             if data_format == "chat":
-                data_format = DataFormat.chat
+                data_format = TrainDataFormat.CHAT
             elif data_format == "completion":
-                data_format = DataFormat.completion
+                data_format = TrainDataFormat.COMPLETION
             else:
                 raise ValueError(
-                    f"String `data_format` must be one of 'chat' or 'completion', but received: {data_format}."
+                    f"String `train_data_format` must be one of 'chat' or 'completion', but received: {data_format}."
                 )
 
         if "train_data_path" not in train_kwargs:
@@ -243,7 +243,7 @@ def finetune(
         return f"databricks/{job.endpoint_name}"
 
     @staticmethod
-    def upload_data(train_data: List[Dict[str, Any]], databricks_unity_catalog_path: str, data_format: DataFormat):
+    def upload_data(train_data: List[Dict[str, Any]], databricks_unity_catalog_path: str, data_format: TrainDataFormat):
         logger.info("Uploading finetuning data to Databricks Unity Catalog...")
         file_path = _save_data_to_local_file(train_data, data_format)
 
@@ -303,7 +303,7 @@ def _create_directory_in_databricks_unity_catalog(w: "WorkspaceClient", databric
         logger.info(f"Successfully created directory {databricks_unity_catalog_path} in Databricks Unity Catalog!")
 
 
-def _save_data_to_local_file(train_data: List[Dict[str, Any]], data_format: DataFormat):
+def _save_data_to_local_file(train_data: List[Dict[str, Any]], data_format: TrainDataFormat):
     import uuid
 
     file_name = f"finetuning_{uuid.uuid4()}.jsonl"
@@ -313,9 +313,9 @@ def _save_data_to_local_file(train_data: List[Dict[str, Any]], data_format: Data
     file_path = os.path.abspath(file_path)
     with open(file_path, "w") as f:
         for item in train_data:
-            if data_format == DataFormat.chat:
+            if data_format == TrainDataFormat.CHAT:
                 _validate_chat_data(item)
-            elif data_format == DataFormat.completion:
+            elif data_format == TrainDataFormat.COMPLETION:
                 _validate_completion_data(item)
 
             f.write(ujson.dumps(item) + "\n")

dspy/clients/lm.py

@@ -20,7 +20,7 @@
 from dspy.adapters.base import Adapter
 from dspy.clients.openai import OpenAIProvider
 from dspy.clients.provider import Provider, TrainingJob
-from dspy.clients.utils_finetune import DataFormat, infer_data_format, validate_data_format
+from dspy.clients.utils_finetune import TrainDataFormat
 from dspy.utils.callback import BaseCallback, with_callbacks
 
 from .base_lm import BaseLM
@@ -46,6 +46,7 @@ def __init__(
         provider=None,
         finetuning_model: Optional[str] = None,
         launch_kwargs: Optional[dict[str, Any]] = None,
+        train_kwargs: Optional[dict[str, Any]] = None,
         **kwargs,
     ):
         """
@@ -79,7 +80,8 @@ def __init__(
         self.callbacks = callbacks or []
         self.num_retries = num_retries
         self.finetuning_model = finetuning_model
-        self.launch_kwargs = launch_kwargs
+        self.launch_kwargs = launch_kwargs or {}
+        self.train_kwargs = train_kwargs or {}
 
         # Handle model-specific configuration for different model families
         model_family = model.split("/")[-1].lower() if "/" in model else model.lower()
@@ -156,18 +158,16 @@ def __call__(self, prompt=None, messages=None, **kwargs):
         return outputs
 
     def launch(self, launch_kwargs: Optional[Dict[str, Any]] = None):
-        launch_kwargs = launch_kwargs or self.launch_kwargs
         self.provider.launch(self, launch_kwargs)
 
     def kill(self, launch_kwargs: Optional[Dict[str, Any]] = None):
-        launch_kwargs = launch_kwargs or self.launch_kwargs
         self.provider.kill(self, launch_kwargs)
 
     def finetune(
         self,
         train_data: List[Dict[str, Any]],
+        train_data_format: Optional[TrainDataFormat],
         train_kwargs: Optional[Dict[str, Any]] = None,
-        data_format: Optional[DataFormat] = None,
     ) -> TrainingJob:
         from dspy import settings as settings
 
@@ -178,27 +178,18 @@ def finetune(
         err = f"Provider {self.provider} does not support fine-tuning."
         assert self.provider.finetunable, err
 
-        # Perform data validation before starting the thread to fail early
-        train_kwargs = train_kwargs or {}
-        if not data_format:
-            adapter = self.infer_adapter()
-            data_format = infer_data_format(adapter)
-        validate_data_format(data=train_data, data_format=data_format)
-
-        # TODO(PR): We can quickly add caching, but doing so requires
-        # adding functions that just call other functions as we had in the last
-        # iteration, unless people have other ideas.
         def thread_function_wrapper():
             return self._run_finetune_job(job)
 
         thread = threading.Thread(target=thread_function_wrapper)
-        model_to_finetune = self.finetuning_model or self.model
+        train_kwargs = train_kwargs or self.train_kwargs
+        model_to_finetune = self.finetuning_model or self.model 
         job = self.provider.TrainingJob(
             thread=thread,
             model=model_to_finetune,
             train_data=train_data,
+            train_data_format=train_data_format,
             train_kwargs=train_kwargs,
-            data_format=data_format,
         )
         thread.start()
 
@@ -212,8 +203,8 @@ def _run_finetune_job(self, job: TrainingJob):
                 job=job,
                 model=job.model,
                 train_data=job.train_data,
+                train_data_format=job.train_data_format,
                 train_kwargs=job.train_kwargs,
-                data_format=job.data_format,
             )
             lm = self.copy(model=model)
             job.set_result(lm)