Upload train_jafari_chatbot.py with huggingface_hub

train_jafari_chatbot.py

@@ -0,0 +1,120 @@
+# /// script
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "transformers>=4.51.0", "datasets", "accelerate", "bitsandbytes"]
+# ///
+
+import sys
+import traceback
+
+print("="*80)
+print("PRODUCTION TRAINING - Jafari Credit WhatsApp Chatbot")
+print("="*80)
+print(f"Python version: {sys.version}")
+
+try:
+    print("\n[1/7] Importing libraries...")
+    from datasets import load_dataset
+    from peft import LoraConfig
+    from trl import SFTTrainer, SFTConfig
+    from transformers import AutoTokenizer
+    import trackio
+    print("✓ All imports successful")
+
+    # Load the dataset from Hub
+    print("\n[2/7] Loading dataset...")
+    dataset = load_dataset("akiliaiafrica/jafari-credit-whatsapp-chatbot", split="train")
+    print(f"✓ Dataset loaded: {len(dataset)} conversations")
+    print(f"Sample keys: {list(dataset[0].keys())}")
+
+    # Create train/eval split
+    print("\n[3/7] Creating train/eval split...")
+    dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
+    print(f"✓ Train: {len(dataset_split['train'])} | Eval: {len(dataset_split['test'])}")
+
+    # Load tokenizer
+    print("\n[4/7] Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    print("✓ Tokenizer loaded")
+
+    # Preprocess dataset
+    print("\n[5/7] Preprocessing dataset with chat template...")
+
+    def convert_to_text(example):
+        text = tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+            add_generation_prompt=False
+        )
+        return {"text": text}
+
+    train_dataset = dataset_split["train"].map(convert_to_text, remove_columns=dataset_split["train"].column_names)
+    eval_dataset = dataset_split["test"].map(convert_to_text, remove_columns=dataset_split["test"].column_names)
+    print(f"✓ Train={len(train_dataset)}, Eval={len(eval_dataset)}")
+    print(f"Sample text length: {len(train_dataset[0]['text'])} chars")
+
+    # LoRA config
+    print("\n[6/7] Configuring LoRA and trainer...")
+    peft_config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    )
+
+    trainer = SFTTrainer(
+        model="Qwen/Qwen3-0.6B",
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        peft_config=peft_config,
+        args=SFTConfig(
+            output_dir="jafari-chatbot-qwen3",
+            num_train_epochs=3,
+            per_device_train_batch_size=2,
+            per_device_eval_batch_size=2,
+            gradient_accumulation_steps=8,
+            gradient_checkpointing=True,
+            learning_rate=2e-4,
+            lr_scheduler_type="cosine",
+            warmup_ratio=0.1,
+            eval_strategy="steps",
+            eval_steps=50,
+            save_strategy="steps",
+            save_steps=50,
+            save_total_limit=3,
+            push_to_hub=True,
+            hub_model_id="akiliaiafrica/jafari-chatbot-qwen3-0.6b",
+            hub_private_repo=True,
+            hub_strategy="every_save",
+            logging_steps=5,
+            report_to="trackio",
+            run_name="jafari-credit-whatsapp-sft",
+            bf16=True,
+            optim="adamw_8bit",
+            max_grad_norm=1.0,
+        )
+    )
+    print("✓ Trainer initialized")
+
+    # Train
+    print("\n[7/7] Starting training...")
+    print("="*80)
+    sys.stdout.flush()
+
+    trainer.train()
+
+    print("\n" + "="*80)
+    print("Training completed! Pushing model to Hub...")
+    trainer.push_to_hub()
+
+    print("\n✓ SUCCESS: Model pushed to akiliaiafrica/jafari-chatbot-qwen3-0.6b")
+    print("="*80)
+
+except Exception as e:
+    print(f"\n\n{'='*80}")
+    print(f"ERROR: {type(e).__name__}: {e}")
+    print("="*80)
+    traceback.print_exc()
+    sys.exit(1)