Upload train_memory_agent.py with huggingface_hub

train_memory_agent.py

@@ -10,22 +10,24 @@
 # ]
 # ///
 
+import os
+import huggingface_hub
 import trackio
 from datasets import load_dataset
 from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
 
+# Explicit login with token from secrets
+token = os.environ.get("HF_TOKEN")
+if token:
+    huggingface_hub.login(token=token)
+    print("Logged in to HF Hub")
+else:
+    print("WARNING: No HF_TOKEN found!")
+
 print("Loading dataset...")
 dataset = load_dataset("erik1988/way2agi-memory-agent-sft", data_files="memory-agent-sft-v3-merged.jsonl", split="train")
 print(f"Dataset loaded: {len(dataset)} examples")
-print(f"Columns: {dataset.column_names}")
-
-# Pre-process: convert messages column from JSON extension to plain format
-# This avoids pyarrow extension type issues with TRL truncation
-def preprocess(example):
-    return example
-
-dataset = dataset.map(preprocess)
 
 dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
 train_dataset = dataset_split["train"]
@@ -46,6 +48,7 @@ config = SFTConfig(
     push_to_hub=True,
     hub_model_id="erik1988/elias-memory-agent-v1",
     hub_strategy="every_save",
+    hub_token=token,
     max_length=None,
     num_train_epochs=5,
     per_device_train_batch_size=2,