fix(cs): initialize score heads and correct training metric reporting

cookbook/client/twinkle/prepare_seq_cls_model.py

@@ -35,10 +35,11 @@
 # The classifier scores the last non-padding token; batched inputs need the pad id to find it.
 text_config.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
 
-torch.manual_seed(42)
-# The load report lists the vision weights as UNEXPECTED and score.weight as MISSING:
-# both are expected, the vision tower is dropped and the new head is initialized here.
+# The load report lists the vision weights as UNEXPECTED and score.weight as MISSING.
 model = AutoModelForSequenceClassification.from_pretrained(model_dir, config=text_config, torch_dtype='auto')
+# Initialize the new head explicitly instead of relying on implicit missing-weight initialization.
+torch.manual_seed(42)
+torch.nn.init.normal_(model.score.weight, mean=0.0, std=text_config.initializer_range)
 model.save_pretrained(args.output_dir)
 tokenizer.save_pretrained(args.output_dir)
 print(f'Saved {type(model).__name__} with {args.num_labels} label(s) to {args.output_dir}')  # noqa: T201

cookbook/client/twinkle/reranker.py

@@ -84,9 +84,10 @@ def train():
 
     # Step 5: training loop. task='reranker' returns one score per pair.
     for step in range(MAX_STEPS):
-        output = model.forward_backward(inputs=batch, task='reranker')
+        model.forward_backward(inputs=batch, task='reranker')
         model.clip_grad_and_step(max_grad_norm=1.0)
-        logger.info(f'Step {step}: loss={output.result["loss"]:.4f}')
+        metric = model.calculate_metric(is_training=True)
+        logger.info(f'Step {step}: {metric.result}')
 
     # Step 6: relevant documents should now score higher, then save the adapter.
     scores = model.forward_only(inputs=batch, task='reranker', return_logits=True).result['logits']

cookbook/client/twinkle/seq_cls.py

@@ -90,9 +90,10 @@ def train():
 
     # Step 5: training loop. task='seq_cls' returns one [num_labels] score row per text.
     for step in range(MAX_STEPS):
-        output = model.forward_backward(inputs=batch, task='seq_cls')
+        model.forward_backward(inputs=batch, task='seq_cls')
         model.clip_grad_and_step(max_grad_norm=1.0)
-        logger.info(f'Step {step}: loss={output.result["loss"]:.4f}')
+        metric = model.calculate_metric(is_training=True)
+        logger.info(f'Step {step}: {metric.result}')
 
     # Step 6: check the predictions, then save the adapter and optimizer state.
     logits = model.forward_only(inputs=batch, task='seq_cls', return_logits=True).result['logits']

src/twinkle/server/model/app.py

@@ -101,6 +101,10 @@ def _make_model_class(base: type) -> type:
 
 def _make_transformers_model(kw: dict[str, Any]) -> Any:
     train_mode = kw.pop('train_mode', 'lora')
+    # The head comes from the persisted config checked by ``_validate_model_task``; HF model
+    # constructors do not accept these deployment keys, unlike Megatron's bridge config.
+    kw.pop('task_type', None)
+    kw.pop('num_labels', None)
     # ``model_cls`` arrives as plain data and is resolved here, so the class object never reaches the
     # constructor -- it would either be rejected by the remote-class layer in safe mode
     # (``TWINKLE_TRUST_REMOTE_CODE=0``) or leak into every worker's constructor kwargs. The schema only

tests/server/integration/task_http_validation.py

@@ -70,11 +70,12 @@ def step(model):
         kwargs = {'task': args.task, 'gradient_accumulation_steps': 2}
         if args.backend == 'megatron':
             kwargs['micro_batch_size'] = 6
-        result = model.forward_backward(inputs=copy.deepcopy(batch), **kwargs).result
-        assert np.isfinite(result['loss'])
+        model.forward_backward(inputs=copy.deepcopy(batch), **kwargs)
         # Call after every backward; the backend gates updates at its GA boundary.
         model.clip_grad_and_step()
-        return float(result['loss'])
+        loss = float(model.calculate_metric(is_training=True).result['loss'])
+        assert np.isfinite(loss)
+        return loss
 
     def save(model, name, samples):
         return model.save(name, save_optimizer=True, consumed_train_samples=samples)

tests/server/model/test_task_head.py

@@ -99,3 +99,13 @@ def tokenizer(text, **kwargs):
     parsed = TypeAdapter(WireInputBatch).validate_python(batch)
     assert len(parsed) == 12
     assert all(type(label) is int for row in parsed for label in row.labels)
+
+
+def test_transformers_backend_does_not_forward_head_keys(monkeypatch):
+    from twinkle.server.model import app
+    from twinkle.server.model.backends import transformers_model
+
+    received = {}
+    monkeypatch.setattr(transformers_model, 'TwinkleCompatTransformersModel', lambda **kwargs: received.update(kwargs))
+    app._make_transformers_model({'model_id': 'unused', 'task_type': 'seq_cls', 'num_labels': 3})
+    assert received == {'model_id': 'unused'}