Fix model loading: use from_pretrained() instead of torch.load() for safetensors format

notebooks/03_ecommerce_finetune.ipynb

@@ -28,7 +28,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# !pip install datasets transformers torch accelerate tokenizers numpy pandas matplotlib scikit-learn wandb huggingface_hub lightgbm"
+    "# !pip install datasets transformers torch accelerate tokenizers numpy pandas matplotlib scikit-learn wandb huggingface_hub lightgbm safetensors"
    ]
   },
   {
@@ -130,11 +130,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Load pre-trained model\n",
-    "config = DomainTransformerConfig.from_preset('24m', vocab_size=hf_tokenizer.vocab_size)\n",
-    "model = DomainTransformerForCausalLM(config)\n",
-    "model.load_state_dict(torch.load('./ecommerce_pretrain_checkpoints/final/model.safetensors',\n",
-    "                                  map_location='cpu', weights_only=True), strict=False)\n",
+    "# Load pre-trained model using from_pretrained (handles safetensors natively)\n",
+    "# Option A: from local checkpoint saved by notebook 02\n",
+    "model = DomainTransformerForCausalLM.from_pretrained('./ecommerce_pretrain_checkpoints/final/')\n",
+    "\n",
+    "# Option B: from HuggingFace Hub (if local not available)\n",
+    "# model = DomainTransformerForCausalLM.from_pretrained('rtferraz/ecommerce-domain-24m')\n",
+    "\n",
     "print(f'Pre-trained model loaded: {sum(p.numel() for p in model.parameters()):,} params')"
    ]
   },
@@ -171,11 +173,9 @@
     "    categories = set(e['category'] for e in events)\n",
     "    n_unique_categories = len(categories)\n",
     "    \n",
-    "    # Temporal features\n",
     "    hours = [e['timestamp'].hour for e in events]\n",
     "    avg_hour = np.mean(hours)\n",
     "    \n",
-    "    # Conversion funnel ratios\n",
     "    cart_rate = n_carts / max(n_views, 1)\n",
     "    purchase_rate = n_purchases / max(n_events, 1)\n",
     "    remove_rate = n_removes / max(n_carts, 1) if n_carts > 0 else 0\n",
@@ -183,16 +183,14 @@
     "    return [\n",
     "        n_events, n_views, n_carts, n_purchases, n_removes,\n",
     "        avg_price, max_price, std_price,\n",
-    "        n_unique_categories,\n",
-    "        avg_hour,\n",
+    "        n_unique_categories, avg_hour,\n",
     "        cart_rate, purchase_rate, remove_rate,\n",
     "    ]\n",
     "\n",
     "FEATURE_NAMES = [\n",
     "    'n_events', 'n_views', 'n_carts', 'n_purchases', 'n_removes',\n",
     "    'avg_price', 'max_price', 'std_price',\n",
-    "    'n_unique_categories',\n",
-    "    'avg_hour',\n",
+    "    'n_unique_categories', 'avg_hour',\n",
     "    'cart_rate', 'purchase_rate', 'remove_rate',\n",
     "]\n",
     "\n",
@@ -201,8 +199,7 @@
     "labels = np.array([1.0 if any(e['event_type'] == 'purchase' for e in seq) else 0.0 for seq in user_sequences])\n",
     "\n",
     "print(f'Features shape: {tabular_features.shape}')\n",
-    "print(f'Labels: {labels.sum():.0f} purchasers / {len(labels)} total ({labels.mean()*100:.1f}%)')\n",
-    "print(f'Feature names: {FEATURE_NAMES}')"
+    "print(f'Labels: {labels.sum():.0f} purchasers / {len(labels)} total ({labels.mean()*100:.1f}%)')"
    ]
   },
   {
@@ -211,7 +208,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Train/test split (80/20, stratified by label)\n",
+    "# Train/test split (80/20, stratified)\n",
     "train_idx, test_idx = train_test_split(\n",
     "    range(len(user_sequences)), test_size=0.2, random_state=42, stratify=labels\n",
     ")\n",
@@ -223,8 +220,8 @@
     "train_labels = labels[train_idx]\n",
     "test_labels = labels[test_idx]\n",
     "\n",
-    "print(f'Train: {len(train_seqs):,} users ({train_labels.mean()*100:.1f}% positive)')\n",
-    "print(f'Test: {len(test_seqs):,} users ({test_labels.mean()*100:.1f}% positive)')"
+    "print(f'Train: {len(train_seqs):,} ({train_labels.mean()*100:.1f}% positive)')\n",
+    "print(f'Test: {len(test_seqs):,} ({test_labels.mean()*100:.1f}% positive)')"
    ]
   },
   {
@@ -257,11 +254,9 @@
     "print(f'  Train AUC: {lgb_train_auc:.4f}')\n",
     "print(f'  Test AUC:  {lgb_test_auc:.4f}')\n",
     "\n",
-    "# Feature importance\n",
     "importance = pd.Series(lgb_model.feature_importances_, index=FEATURE_NAMES).sort_values(ascending=False)\n",
     "print(f'\\nTop features:')\n",
-    "for feat, imp in importance.head(5).items():\n",
-    "    print(f'  {feat}: {imp}')"
+    "for feat, imp in importance.head(5).items(): print(f'  {feat}: {imp}')"
    ]
   },
   {
@@ -270,7 +265,7 @@
    "source": [
     "## Step 4 — JointFusionModel Fine-Tuning\n",
     "\n",
-    "The JointFusionModel combines:\n",
+    "Combines:\n",
     "- **Transaction branch:** Pre-trained DomainTransformer → user embedding\n",
     "- **Tabular branch:** DCNv2 with PLR embeddings on hand-crafted features\n",
     "- **Joint head:** MLP on concatenated embeddings → binary prediction"
@@ -282,21 +277,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Create fine-tuning datasets\n",
-    "MAX_LENGTH = 256  # tokens per user sequence\n",
+    "MAX_LENGTH = 256\n",
     "\n",
     "train_dataset = DomainFinetuneDataset(\n",
-    "    train_seqs, train_features, train_labels,\n",
-    "    builder, hf_tokenizer, max_length=MAX_LENGTH,\n",
-    ")\n",
+    "    train_seqs, train_features, train_labels, builder, hf_tokenizer, max_length=MAX_LENGTH)\n",
     "test_dataset = DomainFinetuneDataset(\n",
-    "    test_seqs, test_features, test_labels,\n",
-    "    builder, hf_tokenizer, max_length=MAX_LENGTH,\n",
-    ")\n",
+    "    test_seqs, test_features, test_labels, builder, hf_tokenizer, max_length=MAX_LENGTH)\n",
     "\n",
-    "print(f'Train dataset: {len(train_dataset)} samples')\n",
-    "print(f'Test dataset: {len(test_dataset)} samples')\n",
-    "print(f'Sample: {set(train_dataset[0].keys())}')"
+    "print(f'Train: {len(train_dataset)}, Test: {len(test_dataset)}')\n",
+    "print(f'Sample keys: {set(train_dataset[0].keys())}')"
    ]
   },
   {
@@ -305,22 +294,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Create JointFusionModel\n",
     "fusion_model = JointFusionModel(\n",
     "    transformer_model=model,\n",
     "    n_tabular_features=len(FEATURE_NAMES),\n",
-    "    n_classes=1,  # binary\n",
-    "    plr_frequencies=32,\n",
-    "    plr_embedding_dim=32,\n",
-    "    dcn_cross_layers=3,\n",
-    "    dcn_deep_layers=2,\n",
-    "    dcn_deep_dim=128,\n",
-    "    head_hidden_dim=128,\n",
-    "    dropout=0.1,\n",
+    "    n_classes=1,\n",
+    "    plr_frequencies=32, plr_embedding_dim=32,\n",
+    "    dcn_cross_layers=3, dcn_deep_layers=2, dcn_deep_dim=128,\n",
+    "    head_hidden_dim=128, dropout=0.1,\n",
     ")\n",
-    "\n",
-    "n_params = sum(p.numel() for p in fusion_model.parameters())\n",
-    "print(f'JointFusion model: {n_params:,} params (transformer + DCNv2 + head)')"
+    "print(f'JointFusion: {sum(p.numel() for p in fusion_model.parameters()):,} params')"
    ]
   },
   {
@@ -348,8 +330,7 @@
     "    logging_steps=20,\n",
     "    eval_steps=100 if USE_GPU else 50,\n",
     "    save_strategy='no',\n",
-    "    bf16=USE_BF16,\n",
-    "    fp16=USE_FP16,\n",
+    "    bf16=USE_BF16, fp16=USE_FP16,\n",
     "    report_to='wandb',\n",
     "    run_name='ecommerce-finetune-joint-5ep',\n",
     "    seed=42,\n",
@@ -369,12 +350,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Get predictions from JointFusion model\n",
     "fusion_model.eval()\n",
     "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
     "fusion_model = fusion_model.to(device)\n",
     "\n",
-    "all_probs, all_labels = [], []\n",
+    "all_probs, all_labels_eval = [], []\n",
     "loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)\n",
     "\n",
     "with torch.no_grad():\n",
@@ -384,12 +364,11 @@
     "        out = fusion_model(**batch)\n",
     "        probs = torch.sigmoid(out['logits'].squeeze(-1))\n",
     "        all_probs.extend(probs.cpu().numpy())\n",
-    "        all_labels.extend(labels_batch.cpu().numpy())\n",
+    "        all_labels_eval.extend(labels_batch.cpu().numpy())\n",
     "\n",
     "all_probs = np.array(all_probs)\n",
-    "all_labels = np.array(all_labels)\n",
-    "\n",
-    "fusion_test_auc = roc_auc_score(all_labels, all_probs)\n",
+    "all_labels_eval = np.array(all_labels_eval)\n",
+    "fusion_test_auc = roc_auc_score(all_labels_eval, all_probs)\n",
     "print(f'JointFusion Test AUC: {fusion_test_auc:.4f}')"
    ]
   },
@@ -399,7 +378,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Comparison table\n",
     "print('=' * 50)\n",
     "print('MODEL COMPARISON — Purchase Prediction (AUC)')\n",
     "print('=' * 50)\n",
@@ -412,7 +390,7 @@
     "    print(f'\\n✅ JointFusion beats LightGBM by {(fusion_test_auc - lgb_test_auc)*100:.2f} percentage points')\n",
     "else:\n",
     "    print(f'\\n⚠️ LightGBM still leads by {(lgb_test_auc - fusion_test_auc)*100:.2f} percentage points')\n",
-    "    print(f'    (Expected with only 3-epoch pre-training. More epochs would improve the transformer embeddings.)')"
+    "    print(f'    (More pre-training epochs and longer context would improve transformer embeddings.)')"
    ]
   },
   {
@@ -421,15 +399,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Loss curve\n",
     "losses = [h['loss'] for h in trainer.state.log_history if 'loss' in h]\n",
     "eval_losses = [h['eval_loss'] for h in trainer.state.log_history if 'eval_loss' in h]\n",
     "\n",
     "fig, ax = plt.subplots(figsize=(10, 5))\n",
     "ax.plot(losses, label='Train Loss', alpha=0.7)\n",
     "if eval_losses:\n",
-    "    eval_steps = np.linspace(0, len(losses), len(eval_losses))\n",
-    "    ax.plot(eval_steps, eval_losses, 'ro-', label='Eval Loss', markersize=4)\n",
+    "    eval_steps_x = np.linspace(0, len(losses), len(eval_losses))\n",
+    "    ax.plot(eval_steps_x, eval_losses, 'ro-', label='Eval Loss', markersize=4)\n",
     "ax.set_xlabel('Step'); ax.set_ylabel('Loss'); ax.set_title('Fine-Tuning Loss')\n",
     "ax.legend(); ax.grid(True, alpha=0.3); plt.tight_layout(); plt.show()"
    ]