Fix poem finetuning example and Tinker compatibility

examples/poem_finetuning_example.py

@@ -6,9 +6,12 @@
 3. Fine-tuning a model using TinkerTrainer
 
 Usage:
+    # Install the optional fine-tuning dependency for real Tinker training
+    poetry install --extras finetune
+
     # Set your API keys
     export OPENAI_API_KEY="your-openai-key"
-    export TINKER_API_KEY="your-tinker-key"  # Optional, runs in mock mode without it
+    export TINKER_API_KEY="your-tinker-key"  # Optional, runs in mock mode without it or without the SDK
 
     # Run the example
     poetry run python examples/poem_finetuning_example.py
@@ -46,10 +49,12 @@ class PoemRequest(BaseModel):
 # =============================================================================
 # Step 2: Create the Curator LLM for data generation
 # =============================================================================
-def create_poem_curator() -> LLM:
-    """Create a Curator LLM configured for poem generation."""
+class PoemCurator(LLM):
+    """Curator LLM configured for poem generation."""
+
+    response_format = Poem
 
-    def prompt_func(row: dict) -> str:
+    def prompt(self, row: dict) -> str:
         """Generate a prompt for creating a poem."""
         return f"""You are a creative poet. Write a {row['style']} poem about "{row['theme']}".
 
@@ -58,7 +63,7 @@ def prompt_func(row: dict) -> str:
 
 After the poem, provide a brief explanation of its meaning and imagery."""
 
-    def parse_func(row: dict, response: Poem) -> dict:
+    def parse(self, row: dict, response: Poem) -> dict:
         """Parse the response into training format."""
         return {
             "theme": response.theme,
@@ -69,12 +74,10 @@ def parse_func(row: dict, response: Poem) -> dict:
             "request_style": row["style"],
         }
 
-    return LLM(
-        model_name="gpt-4o-mini",
-        prompt_func=prompt_func,
-        parse_func=parse_func,
-        response_format=Poem,
-    )
+
+def create_poem_curator() -> LLM:
+    """Create a Curator LLM configured for poem generation."""
+    return PoemCurator(model_name="gpt-4o-mini")
 
 
 # =============================================================================
@@ -119,7 +122,7 @@ def generate_poem_dataset() -> list:
     poem_response = poem_curator(requests)
 
     # Convert response to list
-    poem_data = poem_response.to_list()
+    poem_data = poem_response.dataset.to_list()
 
     print(f"Generated {len(poem_data)} poems")
     return poem_data

src/bespokelabs/curator/finetune/trainer/tinker_trainer.py

@@ -1,5 +1,6 @@
 """TinkerTrainer implementation for LoRA fine-tuning via Tinker API."""
 
+import inspect
 import time
 from typing import Any, Dict, List, Optional
 
@@ -17,10 +18,16 @@
 )
 from bespokelabs.curator.log import logger
 
+tinker = None
 try:
-    import tinker
-
-    TINKER_AVAILABLE = True
+    import tinker as _tinker
+
+    required_tinker_attrs = ("ServiceClient", "Datum", "ModelInput", "EncodedTextChunk")
+    if all(hasattr(_tinker, attr) for attr in required_tinker_attrs):
+        tinker = _tinker
+        TINKER_AVAILABLE = True
+    else:
+        TINKER_AVAILABLE = False
 except ImportError:
     TINKER_AVAILABLE = False
 
@@ -88,16 +95,26 @@ def _initialize_client(self) -> None:
 
         try:
             self._service_client = tinker.ServiceClient(api_key=self.config.api_key)
-            self._training_client = self._service_client.create_lora_training_client(
-                base_model=self.config.base_model,
-                rank=self.config.lora_config.rank,
-                alpha=self.config.lora_config.alpha,
-                dropout=self.config.lora_config.dropout,
-                target_modules=self.config.lora_config.target_modules,
-                train_mlp=True,
-                train_attn=True,
-                train_unembed=True,
+            training_client_kwargs = {
+                "base_model": self.config.base_model,
+                "rank": self.config.lora_config.rank,
+                "alpha": self.config.lora_config.alpha,
+                "dropout": self.config.lora_config.dropout,
+                "target_modules": self.config.lora_config.target_modules,
+                "train_mlp": True,
+                "train_attn": True,
+                "train_unembed": True,
+            }
+            supported_kwargs, ignored_kwargs = self._get_supported_kwargs(
+                self._service_client.create_lora_training_client,
+                training_client_kwargs,
             )
+            if ignored_kwargs:
+                logger.warning(
+                    "Installed Tinker SDK does not support LoRA args %s; continuing with supported args only.",
+                    ", ".join(ignored_kwargs),
+                )
+            self._training_client = self._service_client.create_lora_training_client(**supported_kwargs)
             self._tokenizer = self._training_client.get_tokenizer()
 
             logger.info(f"TinkerTrainer initialized for model: {self.config.base_model}")
@@ -109,6 +126,21 @@ def _initialize_client(self) -> None:
             self._training_client = None
             self._tokenizer = None
 
+    @staticmethod
+    def _get_supported_kwargs(callable_obj: Any, kwargs: Dict[str, Any]) -> tuple[Dict[str, Any], List[str]]:
+        """Filter keyword arguments to those accepted by the callable."""
+        try:
+            parameters = inspect.signature(callable_obj).parameters
+        except (TypeError, ValueError):
+            return kwargs, []
+
+        if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()):
+            return kwargs, []
+
+        supported_kwargs = {key: value for key, value in kwargs.items() if key in parameters}
+        ignored_kwargs = [key for key in kwargs if key not in supported_kwargs]
+        return supported_kwargs, ignored_kwargs
+
     def save_checkpoint(self, name: str, step: int, epoch: int, loss: float) -> Optional[CheckpointInfo]:
         """Save a training checkpoint.
 
@@ -488,6 +520,11 @@ def _get_effective_target_weight(batch_data: List[Any]) -> float:
             if weights is None:
                 continue
 
+            if hasattr(weights, "data"):
+                weights = weights.data
+            elif isinstance(weights, dict) and "data" in weights:
+                weights = weights["data"]
+
             found_weights = True
             effective_weight += sum(float(weight) for weight in weights)
 

tests/finetune/test_trainer.py

@@ -151,6 +151,77 @@ def test_get_sampling_client(self, trainer):
         # In mock mode, client is None
         assert client is None
 
+    def test_initialize_client_filters_unsupported_lora_args(self, monkeypatch):
+        """Test client initialization against SDKs with narrower LoRA signatures."""
+        captured = {}
+
+        class FakeTrainingClient:
+            def get_tokenizer(self):
+                return SimpleNamespace(vocab_size=123)
+
+        class FakeServiceClient:
+            def __init__(self, api_key):
+                captured["api_key"] = api_key
+
+            def create_lora_training_client(self, base_model, rank, train_mlp=True, train_attn=True, train_unembed=True):
+                captured["kwargs"] = {
+                    "base_model": base_model,
+                    "rank": rank,
+                    "train_mlp": train_mlp,
+                    "train_attn": train_attn,
+                    "train_unembed": train_unembed,
+                }
+                return FakeTrainingClient()
+
+        monkeypatch.setattr(tinker_trainer_module, "TINKER_AVAILABLE", True)
+        monkeypatch.setattr(tinker_trainer_module, "tinker", SimpleNamespace(ServiceClient=FakeServiceClient))
+
+        trainer = TinkerTrainer(TinkerTrainerConfig(base_model="Qwen3-8B", epochs=1, api_key="test-key"))
+
+        assert captured["api_key"] == "test-key"
+        assert captured["kwargs"] == {
+            "base_model": "Qwen3-8B",
+            "rank": 16,
+            "train_mlp": True,
+            "train_attn": True,
+            "train_unembed": True,
+        }
+        assert trainer._tokenizer.vocab_size == 123
+
+    def test_initialize_client_keeps_kwargs_for_var_keyword_signatures(self, monkeypatch):
+        """Test client initialization preserves kwargs when the SDK accepts **kwargs."""
+        captured = {}
+
+        class FakeTrainingClient:
+            def get_tokenizer(self):
+                return SimpleNamespace(vocab_size=123)
+
+        class FakeServiceClient:
+            def __init__(self, api_key):
+                captured["api_key"] = api_key
+
+            def create_lora_training_client(self, **kwargs):
+                captured["kwargs"] = kwargs
+                return FakeTrainingClient()
+
+        monkeypatch.setattr(tinker_trainer_module, "TINKER_AVAILABLE", True)
+        monkeypatch.setattr(tinker_trainer_module, "tinker", SimpleNamespace(ServiceClient=FakeServiceClient))
+
+        trainer = TinkerTrainer(TinkerTrainerConfig(base_model="Qwen3-8B", epochs=1, api_key="test-key"))
+
+        assert captured["api_key"] == "test-key"
+        assert captured["kwargs"] == {
+            "base_model": "Qwen3-8B",
+            "rank": 16,
+            "alpha": 32,
+            "dropout": 0.05,
+            "target_modules": ["q_proj", "v_proj", "k_proj", "o_proj"],
+            "train_mlp": True,
+            "train_attn": True,
+            "train_unembed": True,
+        }
+        assert trainer._tokenizer.vocab_size == 123
+
     def test_training_step_normalizes_loss_by_weighted_targets(self, trainer, monkeypatch):
         """Test loss uses active target weights, not total prompt length."""
 
@@ -185,6 +256,25 @@ def __init__(self, input_tokens, weights):
         assert loss == pytest.approx(1.0)
         assert tokens_processed == 4
 
+    def test_effective_target_weight_handles_tensor_wrapped_weights(self):
+        """Test weight extraction from tensor-like wrappers used by Tinker."""
+
+        class FakeTensorData:
+            def __init__(self, data):
+                self.data = data
+
+            def __iter__(self):
+                return iter([("data", self.data)])
+
+        datum = SimpleNamespace(
+            model_input=SimpleNamespace(chunks=[SimpleNamespace(tokens=[10, 11, 12])]),
+            loss_fn_inputs={"weights": FakeTensorData([0.0, 0.5, 1.0])},
+        )
+
+        effective_weight = TinkerTrainer._get_effective_target_weight([datum])
+
+        assert effective_weight == pytest.approx(1.5)
+
 
 class TestCustomTrainer:
     """Tests for custom trainer implementations."""