Update app.py

app.py

@@ -1,9 +1,6 @@
 """
-AI Python Code Model Trainer
-Hugging Face Space for continuous training with auto-resume
-Username: himu1780 | Model: ai-python-model
-
-FINAL VERSION - All optimizations applied
+AI Python Code Model Trainer - FIXED VERSION
+Dataset: jtatman/python-code-dataset-500k
 """
 
 import os
@@ -12,7 +9,7 @@ import gradio as gr
 import threading
 import time
 from datetime import datetime
-from huggingface_hub import HfApi, login
+from huggingface_hub import HfApi, login, create_repo
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -22,7 +19,6 @@ from transformers import (
 )
 from datasets import load_dataset, Dataset
 
-# Try to import torch for memory cleanup
 try:
     import torch
     TORCH_AVAILABLE = True
@@ -35,7 +31,6 @@ MODEL_REPO = f"{HF_USERNAME}/ai-python-model"
 DATASET_NAME = "jtatman/python-code-dataset-500k"
 BASE_MODEL = "gpt2"
 
-# Training hyperparameters (Memory optimized)
 BATCH_SIZE = 1
 GRADIENT_ACCUMULATION = 8
 SAVE_STEPS = 500
@@ -45,9 +40,8 @@ LEARNING_RATE = 5e-5
 MAX_STEPS_PER_SESSION = 10000
 EXAMPLES_PER_SESSION = 50000
 
-# Continuous training settings
-CONTINUOUS_TRAINING = True  # Set False to stop after one session
-WAIT_BETWEEN_SESSIONS = 60  # Seconds to wait before next session
+CONTINUOUS_TRAINING = True
+WAIT_BETWEEN_SESSIONS = 60
 
 # ============ GLOBAL STATE ============
 training_status = {
@@ -58,13 +52,13 @@ training_status = {
     "start_time": None,
     "message": "Initializing...",
     "session_count": 0,
+    "last_error": "",
 }
 
 stop_requested = False
 
 # ============ MEMORY CLEANUP ============
 def cleanup_memory():
-    """Free up memory after training"""
     gc.collect()
     if TORCH_AVAILABLE and torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -72,19 +66,26 @@ def cleanup_memory():
 
 # ============ AUTHENTICATION ============
 def authenticate():
-    """Login to Hugging Face Hub"""
     token = os.environ.get("HF_TOKEN")
     if token:
         login(token=token)
         training_status["message"] = "✅ Authenticated with Hugging Face"
+        
+        try:
+            api = HfApi()
+            api.create_repo(repo_id=MODEL_REPO, exist_ok=True)
+            print(f"[INFO] Repo {MODEL_REPO} ready")
+        except Exception as e:
+            print(f"[WARN] Repo check: {e}")
+        
         return True
     else:
         training_status["message"] = "❌ HF_TOKEN not found in secrets!"
+        training_status["last_error"] = "Add HF_TOKEN to Space secrets"
         return False
 
 # ============ MODEL LOADING ============
 def load_model_and_tokenizer():
-    """Load model from Hub (resume) or start fresh from base model"""
     global training_status
     
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
@@ -92,7 +93,7 @@ def load_model_and_tokenizer():
     
     try:
         training_status["message"] = f"🔄 Attempting to resume from {MODEL_REPO}..."
-        model = AutoModelForCausalLM.from_pretrained(MODEL_REPO)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, trust_remote_code=True)
         training_status["message"] = f"✅ Resumed from {MODEL_REPO}"
         print(f"[INFO] Resumed training from {MODEL_REPO}")
     except Exception as e:
@@ -104,68 +105,92 @@ def load_model_and_tokenizer():
 
 # ============ DATASET PROCESSING ============
 def prepare_dataset(tokenizer):
-    """Load and prepare dataset"""
     global training_status
-    training_status["message"] = "📥 Loading dataset (streaming mode)..."
+    training_status["message"] = "📥 Loading dataset..."
     
     try:
+        # Load dataset in streaming mode
         dataset = load_dataset(DATASET_NAME, split="train", streaming=True)
+        
+        # Take only what we need
         dataset = dataset.take(EXAMPLES_PER_SESSION)
         
-        def tokenize_function(examples):
-            texts = []
-            instructions = examples.get("instruction", [])
-            outputs = examples.get("output", [])
-            
-            for instruction, output in zip(instructions, outputs):
-                if instruction and output:
-                    text = f"### Instruction:\n{instruction}\n\n### Response:\n{output}"
-                    texts.append(text)
-            
-            if not texts:
-                texts = [""]
-            
-            result = tokenizer(
-                texts,
-                truncation=True,
-                max_length=MAX_LENGTH,
-                padding="max_length",
-                return_tensors=None,
-            )
-            result["labels"] = result["input_ids"].copy()
-            return result
+        training_status["message"] = "🔄 Processing examples..."
         
-        tokenized_dataset = dataset.map(
-            tokenize_function,
-            batched=True,
-            batch_size=100,
-            remove_columns=["instruction", "output"],
-        )
+        all_examples = []
+        count = 0
         
-        training_status["message"] = "🔄 Converting dataset for Trainer..."
+        for example in dataset:
+            try:
+                # Get instruction and output from dataset
+                # This dataset has: instruction, output, system columns
+                instruction = example.get("instruction", "")
+                output = example.get("output", "")
+                
+                # Skip if empty
+                if not instruction or not output:
+                    continue
+                
+                # Make sure they are strings
+                if not isinstance(instruction, str):
+                    instruction = str(instruction)
+                if not isinstance(output, str):
+                    output = str(output)
+                
+                # Create training text
+                text = f"### Instruction:\n{instruction}\n\n### Response:\n{output}"
+                
+                # Tokenize
+                tokenized = tokenizer(
+                    text,
+                    truncation=True,
+                    max_length=MAX_LENGTH,
+                    padding="max_length",
+                    return_tensors=None,
+                )
+                
+                # Create example with only needed fields
+                processed_example = {
+                    "input_ids": tokenized["input_ids"],
+                    "attention_mask": tokenized["attention_mask"],
+                    "labels": tokenized["input_ids"].copy(),
+                }
+                
+                all_examples.append(processed_example)
+                count += 1
+                
+                # Progress update
+                if count % 5000 == 0:
+                    training_status["message"] = f"📥 Processed {count:,}/{EXAMPLES_PER_SESSION:,} examples..."
+                    print(f"[INFO] Processed {count:,} examples...")
+                
+                if count >= EXAMPLES_PER_SESSION:
+                    break
+                    
+            except Exception as e:
+                # Skip problematic examples
+                continue
         
-        all_examples = []
-        for i, example in enumerate(tokenized_dataset):
-            all_examples.append(example)
-            if i % 5000 == 0:
-                training_status["message"] = f"📥 Loaded {i:,}/{EXAMPLES_PER_SESSION:,} examples..."
-            if i >= EXAMPLES_PER_SESSION - 1:
-                break
+        if len(all_examples) == 0:
+            raise ValueError("No valid examples found in dataset!")
         
+        # Create HuggingFace Dataset
         train_dataset = Dataset.from_list(all_examples)
         
         training_status["message"] = f"✅ Dataset ready: {len(train_dataset):,} examples"
+        print(f"[INFO] Dataset ready: {len(train_dataset):,} examples")
         return train_dataset
         
     except Exception as e:
         training_status["message"] = f"❌ Dataset error: {str(e)}"
+        training_status["last_error"] = str(e)
         print(f"[ERROR] Dataset preparation failed: {e}")
+        import traceback
+        traceback.print_exc()
         raise e
 
 # ============ CUSTOM TRAINER ============
 class StatusTrainer(Trainer):
-    """Custom trainer with status updates and stop support"""
-    
     def training_step(self, model, inputs):
         global stop_requested
         if stop_requested:
@@ -184,7 +209,6 @@ class StatusTrainer(Trainer):
 
 # ============ SINGLE TRAINING SESSION ============
 def run_training_session():
-    """Run a single training session"""
     global training_status, stop_requested
     
     model = None
@@ -197,6 +221,10 @@ def run_training_session():
         model, tokenizer = load_model_and_tokenizer()
         train_dataset = prepare_dataset(tokenizer)
         
+        if len(train_dataset) == 0:
+            training_status["message"] = "❌ Empty dataset!"
+            return False
+        
         data_collator = DataCollatorForLanguageModeling(
             tokenizer=tokenizer,
             mlm=False,
@@ -220,7 +248,7 @@ def run_training_session():
             max_steps=MAX_STEPS_PER_SESSION,
             fp16=False,
             dataloader_num_workers=0,
-            remove_unused_columns=False,
+            remove_unused_columns=True,
         )
         
         trainer = StatusTrainer(
@@ -228,11 +256,14 @@ def run_training_session():
             args=training_args,
             train_dataset=train_dataset,
             data_collator=data_collator,
-            tokenizer=tokenizer,
+            processing_class=tokenizer,
         )
         
         training_status["message"] = "🏃 Training in progress..."
+        print("[INFO] Starting training...")
         trainer.train()
+        
+        print("[INFO] Pushing to hub...")
         trainer.push_to_hub()
         
         training_status["session_count"] += 1
@@ -243,18 +274,21 @@ def run_training_session():
         training_status["message"] = "⏹️ Training stopped by user"
         return False
     except Exception as e:
-        training_status["message"] = f"❌ Error: {str(e)}"
+        training_status["message"] = f"❌ Error: {str(e)[:100]}"
+        training_status["last_error"] = str(e)
         print(f"[ERROR] Training failed: {e}")
         import traceback
         traceback.print_exc()
         return False
     finally:
-        del model, trainer
+        if model is not None:
+            del model
+        if trainer is not None:
+            del trainer
         cleanup_memory()
 
 # ============ MAIN TRAINING LOOP ============
 def start_training():
-    """Main training function with continuous loop"""
     global training_status, stop_requested
     
     if training_status["is_training"]:
@@ -289,7 +323,6 @@ def start_training():
 
 # ============ GRADIO INTERFACE ============
 def get_status():
-    """Get current training status"""
     elapsed = ""
     if training_status["start_time"]:
         delta = datetime.now() - training_status["start_time"]
@@ -306,6 +339,7 @@ def get_status():
     continuous_str = "✅ Enabled" if CONTINUOUS_TRAINING else "❌ Disabled"
     elapsed_str = elapsed if elapsed else "N/A"
     effective_batch = BATCH_SIZE * GRADIENT_ACCUMULATION
+    error_str = training_status["last_error"][:100] if training_status["last_error"] else "None"
     
     return f"""
 ## 🤖 AI Python Model Trainer
@@ -316,6 +350,7 @@ def get_status():
 | **State** | {state_str} |
 | **Message** | {training_status["message"]} |
 | **Sessions Completed** | {training_status["session_count"]} |
+| **Last Error** | {error_str} |
 
 ### Progress
 | Metric | Value |
@@ -335,7 +370,6 @@ def get_status():
 """
 
 def start_training_async():
-    """Start training in background"""
     if training_status["is_training"]:
         return "⚠️ Training already in progress!"
     thread = threading.Thread(target=start_training, daemon=True)
@@ -343,7 +377,6 @@ def start_training_async():
     return "🚀 Training started in background!"
 
 def stop_training():
-    """Stop training"""
     global stop_requested
     if not training_status["is_training"]:
         return "⚠️ No training in progress"
@@ -353,7 +386,6 @@ def stop_training():
 
 # ============ AUTO-START ============
 def auto_start():
-    """Auto-start continuous training on Space launch"""
     time.sleep(10)
     while True:
         if not training_status["is_training"] and not stop_requested: