Update app.py

app.py

@@ -1,380 +1,392 @@
-"""
-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
-"""
-
-import os
-import gc
-import gradio as gr
-import threading
-import time
-from datetime import datetime
-from huggingface_hub import HfApi, login
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TrainingArguments,
-    Trainer,
-    DataCollatorForLanguageModeling,
-)
-from datasets import load_dataset, Dataset
-
-# Try to import torch for memory cleanup
-try:
-    import torch
-    TORCH_AVAILABLE = True
-except ImportError:
-    TORCH_AVAILABLE = False
-
-# ============ CONFIGURATION ============
-HF_USERNAME = "himu1780"
-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
-LOGGING_STEPS = 50
-MAX_LENGTH = 256
-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
-
-# ============ GLOBAL STATE ============
-training_status = {
-    "is_training": False,
-    "current_step": 0,
-    "total_loss": 0,
-    "last_save": "Never",
-    "start_time": None,
-    "message": "Initializing...",
-    "session_count": 0,
-}
-
-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()
-    print("[INFO] Memory cleaned up")
-
-# ============ 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"
-        return True
-    else:
-        training_status["message"] = "❌ HF_TOKEN not found in 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)
-    tokenizer.pad_token = tokenizer.eos_token
-    
-    try:
-        training_status["message"] = f"🔄 Attempting to resume from {MODEL_REPO}..."
-        model = AutoModelForCausalLM.from_pretrained(MODEL_REPO)
-        training_status["message"] = f"✅ Resumed from {MODEL_REPO}"
-        print(f"[INFO] Resumed training from {MODEL_REPO}")
-    except Exception as e:
-        training_status["message"] = f"🆕 Starting fresh from {BASE_MODEL}"
-        model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
-        print(f"[INFO] Starting fresh from {BASE_MODEL}: {e}")
-    
-    return model, tokenizer
-
-# ============ DATASET PROCESSING ============
-def prepare_dataset(tokenizer):
-    """Load and prepare dataset"""
-    global training_status
-    training_status["message"] = "📥 Loading dataset (streaming mode)..."
-    
-    try:
-        dataset = load_dataset(DATASET_NAME, split="train", streaming=True)
-        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
-        
-        tokenized_dataset = dataset.map(
-            tokenize_function,
-            batched=True,
-            batch_size=100,
-            remove_columns=["instruction", "output"],
-        )
-        
-        training_status["message"] = "🔄 Converting dataset for Trainer..."
-        
-        all_examples = []
-        for i, example in enumerate(tokenized_dataset):
-            all_examples.append(example)
-            # Progress every 5000 (IMPROVED)
-            if i % 5000 == 0:
-                training_status["message"] = f"📥 Loaded {i:,}/{EXAMPLES_PER_SESSION:,} examples..."
-            if i >= EXAMPLES_PER_SESSION - 1:
-                break
-        
-        train_dataset = Dataset.from_list(all_examples)
-        
-        training_status["message"] = f"✅ Dataset ready: {len(train_dataset):,} examples"
-        return train_dataset
-        
-    except Exception as e:
-        training_status["message"] = f"❌ Dataset error: {str(e)}"
-        print(f"[ERROR] Dataset preparation failed: {e}")
-        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:
-            raise KeyboardInterrupt("Stop requested by user")
-        return super().training_step(model, inputs)
-    
-    def log(self, logs):
-        super().log(logs)
-        if "loss" in logs:
-            training_status["total_loss"] = logs["loss"]
-        training_status["current_step"] = self.state.global_step
-    
-    def save_model(self, output_dir=None, _internal_call=False):
-        super().save_model(output_dir, _internal_call)
-        training_status["last_save"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-
-# ============ SINGLE TRAINING SESSION ============
-def run_training_session():
-    """Run a single training session"""
-    global training_status, stop_requested
-    
-    model = None
-    trainer = None
-    
-    try:
-        if not authenticate():
-            return False
-        
-        model, tokenizer = load_model_and_tokenizer()
-        train_dataset = prepare_dataset(tokenizer)
-        
-        data_collator = DataCollatorForLanguageModeling(
-            tokenizer=tokenizer,
-            mlm=False,
-        )
-        
-        training_args = TrainingArguments(
-            output_dir="./temp_checkpoints",
-            overwrite_output_dir=True,
-            per_device_train_batch_size=BATCH_SIZE,
-            gradient_accumulation_steps=GRADIENT_ACCUMULATION,
-            learning_rate=LEARNING_RATE,
-            warmup_steps=100,
-            weight_decay=0.01,
-            logging_steps=LOGGING_STEPS,
-            save_steps=SAVE_STEPS,
-            save_total_limit=1,
-            push_to_hub=True,
-            hub_model_id=MODEL_REPO,
-            hub_strategy="every_save",
-            report_to="none",
-            max_steps=MAX_STEPS_PER_SESSION,
-            fp16=False,
-            dataloader_num_workers=0,
-            remove_unused_columns=False,
-        )
-        
-        trainer = StatusTrainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_dataset,
-            data_collator=data_collator,
-            tokenizer=tokenizer,
-        )
-        
-        training_status["message"] = "🏃 Training in progress..."
-        trainer.train()
-        trainer.push_to_hub()
-        
-        training_status["session_count"] += 1
-        training_status["message"] = f"✅ Session {training_status['session_count']} completed!"
-        return True
-        
-    except KeyboardInterrupt:
-        training_status["message"] = "⏹️ Training stopped by user"
-        return False
-    except Exception as e:
-        training_status["message"] = f"❌ Error: {str(e)}"
-        print(f"[ERROR] Training failed: {e}")
-        import traceback
-        traceback.print_exc()
-        return False
-    finally:
-        # MEMORY CLEANUP (IMPROVED)
-        del model, 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"]:
-        return "Training already in progress!"
-    
-    training_status["is_training"] = True
-    training_status["start_time"] = datetime.now()
-    stop_requested = False
-    
-    # CONTINUOUS TRAINING LOOP (IMPROVED)
-    while not stop_requested:
-        training_status["message"] = f"🚀 Starting session {training_status['session_count'] + 1}..."
-        
-        success = run_training_session()
-        
-        if stop_requested:
-            break
-        
-        if not CONTINUOUS_TRAINING:
-            break
-        
-        if success:
-            training_status["message"] = f"⏳ Waiting {WAIT_BETWEEN_SESSIONS}s before next session..."
-            time.sleep(WAIT_BETWEEN_SESSIONS)
-        else:
-            training_status["message"] = "⚠️ Session failed, retrying in 60s..."
-            time.sleep(60)
-    
-    training_status["is_training"] = False
-    stop_requested = False
-    training_status["message"] = f"✅ Training finished! Total sessions: {training_status['session_count']}"
-    return training_status["message"]
-
-# ============ GRADIO INTERFACE ============
-def get_status():
-    """Get current training status"""
-    elapsed = ""
-    if training_status["start_time"]:
-        delta = datetime.now() - training_status["start_time"]
-        hours, remainder = divmod(int(delta.total_seconds()), 3600)
-        minutes, seconds = divmod(remainder, 60)
-        elapsed = f"{hours}h {minutes}m {seconds}s"
-    
-    return f"""
-## 🤖 AI Python Model Trainer
-
-### Status
-| Item | Value |
-|------|-------|
-| **State** | {"🟢 Training" if training_status["is_training"] else "🔴 Stopped"} |
-| **Message** | {training_status["message"]} |
-| **Sessions Completed** | {training_status["session_count"]} |
-
-### Progress
-| Metric | Value |
-|--------|-------|
-| **Current Step** | {training_status["current_step"]:,} / {MAX_STEPS_PER_SESSION:,} |
-| **Current Loss** | {training_status["total_loss"]:.4f if training_status["total_loss"] else "N/A"} |
-| **Last Checkpoint** | {training_status["last_save"]} |
-| **Elapsed Time** | {elapsed if elapsed else "N/A"} |
-
-### Configuration
-| Setting | Value |
-|---------|-------|
-| **Model Repo** | [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO}) |
-| **Continuous Mode** | {"✅ Enabled" if CONTINUOUS_TRAINING else "❌ Disabled"} |
-| **Batch Size** | {BATCH_SIZE} (effective: {BATCH_SIZE * GRADIENT_ACCUMULATION}) |
-| **Max Steps/Session** | {MAX_STEPS_PER_SESSION:,} |
-"""
-
-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)
-    thread.start()
-    return "🚀 Training started in background!"
-
-def stop_training():
-    """Stop training"""
-    global stop_requested
-    if not training_status["is_training"]:
-        return "⚠️ No training in progress"
-    stop_requested = True
-    training_status["message"] = "⏹️ Stopping after current step..."
-    return "⏹️ Stop requested"
-
-# ============ 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:
-            print("[INFO] Auto-starting training session...")
-            start_training()
-        time.sleep(WAIT_BETWEEN_SESSIONS)
-
-auto_thread = threading.Thread(target=auto_start, daemon=True)
-auto_thread.start()
-
-# ============ GRADIO APP ============
-with gr.Blocks(title="AI Python Trainer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🐍 AI Python Code Model Trainer")
-    gr.Markdown(f"**Continuous training** on `{DATASET_NAME}` with auto-checkpoint")
-    
-    status_display = gr.Markdown(get_status)
-    
-    with gr.Row():
-        start_btn = gr.Button("▶️ Start Training", variant="primary")
-        stop_btn = gr.Button("⏹️ Stop Training", variant="stop")
-        refresh_btn = gr.Button("🔄 Refresh Status")
-    
-    output = gr.Textbox(label="Output", interactive=False)
-    
-    start_btn.click(start_training_async, outputs=output)
-    stop_btn.click(stop_training, outputs=output)
-    refresh_btn.click(get_status, outputs=status_display)
-    
-    demo.load(get_status, outputs=status_display, every=30)
-
-demo.launch()
+"""
+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
+"""
+
+import os
+import gc
+import gradio as gr
+import threading
+import time
+from datetime import datetime
+from huggingface_hub import HfApi, login
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+)
+from datasets import load_dataset, Dataset
+
+# Try to import torch for memory cleanup
+try:
+    import torch
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+
+# ============ CONFIGURATION ============
+HF_USERNAME = "himu1780"
+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
+LOGGING_STEPS = 50
+MAX_LENGTH = 256
+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
+
+# ============ GLOBAL STATE ============
+training_status = {
+    "is_training": False,
+    "current_step": 0,
+    "total_loss": 0,
+    "last_save": "Never",
+    "start_time": None,
+    "message": "Initializing...",
+    "session_count": 0,
+}
+
+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()
+    print("[INFO] Memory cleaned up")
+
+# ============ 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"
+        return True
+    else:
+        training_status["message"] = "❌ HF_TOKEN not found in 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)
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    try:
+        training_status["message"] = f"🔄 Attempting to resume from {MODEL_REPO}..."
+        model = AutoModelForCausalLM.from_pretrained(MODEL_REPO)
+        training_status["message"] = f"✅ Resumed from {MODEL_REPO}"
+        print(f"[INFO] Resumed training from {MODEL_REPO}")
+    except Exception as e:
+        training_status["message"] = f"🆕 Starting fresh from {BASE_MODEL}"
+        model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
+        print(f"[INFO] Starting fresh from {BASE_MODEL}: {e}")
+    
+    return model, tokenizer
+
+# ============ DATASET PROCESSING ============
+def prepare_dataset(tokenizer):
+    """Load and prepare dataset"""
+    global training_status
+    training_status["message"] = "📥 Loading dataset (streaming mode)..."
+    
+    try:
+        dataset = load_dataset(DATASET_NAME, split="train", streaming=True)
+        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
+        
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            batch_size=100,
+            remove_columns=["instruction", "output"],
+        )
+        
+        training_status["message"] = "🔄 Converting dataset for Trainer..."
+        
+        all_examples = []
+        for i, example in enumerate(tokenized_dataset):
+            all_examples.append(example)
+            # Progress every 5000 (IMPROVED)
+            if i % 5000 == 0:
+                training_status["message"] = f"📥 Loaded {i:,}/{EXAMPLES_PER_SESSION:,} examples..."
+            if i >= EXAMPLES_PER_SESSION - 1:
+                break
+        
+        train_dataset = Dataset.from_list(all_examples)
+        
+        training_status["message"] = f"✅ Dataset ready: {len(train_dataset):,} examples"
+        return train_dataset
+        
+    except Exception as e:
+        training_status["message"] = f"❌ Dataset error: {str(e)}"
+        print(f"[ERROR] Dataset preparation failed: {e}")
+        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:
+            raise KeyboardInterrupt("Stop requested by user")
+        return super().training_step(model, inputs)
+    
+    def log(self, logs):
+        super().log(logs)
+        if "loss" in logs:
+            training_status["total_loss"] = logs["loss"]
+        training_status["current_step"] = self.state.global_step
+    
+    def save_model(self, output_dir=None, _internal_call=False):
+        super().save_model(output_dir, _internal_call)
+        training_status["last_save"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+# ============ SINGLE TRAINING SESSION ============
+def run_training_session():
+    """Run a single training session"""
+    global training_status, stop_requested
+    
+    model = None
+    trainer = None
+    
+    try:
+        if not authenticate():
+            return False
+        
+        model, tokenizer = load_model_and_tokenizer()
+        train_dataset = prepare_dataset(tokenizer)
+        
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False,
+        )
+        
+        training_args = TrainingArguments(
+            output_dir="./temp_checkpoints",
+            overwrite_output_dir=True,
+            per_device_train_batch_size=BATCH_SIZE,
+            gradient_accumulation_steps=GRADIENT_ACCUMULATION,
+            learning_rate=LEARNING_RATE,
+            warmup_steps=100,
+            weight_decay=0.01,
+            logging_steps=LOGGING_STEPS,
+            save_steps=SAVE_STEPS,
+            save_total_limit=1,
+            push_to_hub=True,
+            hub_model_id=MODEL_REPO,
+            hub_strategy="every_save",
+            report_to="none",
+            max_steps=MAX_STEPS_PER_SESSION,
+            fp16=False,
+            dataloader_num_workers=0,
+            remove_unused_columns=False,
+        )
+        
+        trainer = StatusTrainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            data_collator=data_collator,
+            tokenizer=tokenizer,
+        )
+        
+        training_status["message"] = "🏃 Training in progress..."
+        trainer.train()
+        trainer.push_to_hub()
+        
+        training_status["session_count"] += 1
+        training_status["message"] = f"✅ Session {training_status['session_count']} completed!"
+        return True
+        
+    except KeyboardInterrupt:
+        training_status["message"] = "⏹️ Training stopped by user"
+        return False
+    except Exception as e:
+        training_status["message"] = f"❌ Error: {str(e)}"
+        print(f"[ERROR] Training failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+    finally:
+        # MEMORY CLEANUP (IMPROVED)
+        del model, 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"]:
+        return "Training already in progress!"
+    
+    training_status["is_training"] = True
+    training_status["start_time"] = datetime.now()
+    stop_requested = False
+    
+    # CONTINUOUS TRAINING LOOP (IMPROVED)
+    while not stop_requested:
+        training_status["message"] = f"🚀 Starting session {training_status['session_count'] + 1}..."
+        
+        success = run_training_session()
+        
+        if stop_requested:
+            break
+        
+        if not CONTINUOUS_TRAINING:
+            break
+        
+        if success:
+            training_status["message"] = f"⏳ Waiting {WAIT_BETWEEN_SESSIONS}s before next session..."
+            time.sleep(WAIT_BETWEEN_SESSIONS)
+        else:
+            training_status["message"] = "⚠️ Session failed, retrying in 60s..."
+            time.sleep(60)
+    
+    training_status["is_training"] = False
+    stop_requested = False
+    training_status["message"] = f"✅ Training finished! Total sessions: {training_status['session_count']}"
+    return training_status["message"]
+
+# ============ GRADIO INTERFACE ============
+def get_status():
+    """Get current training status"""
+    elapsed = ""
+    if training_status["start_time"]:
+        delta = datetime.now() - training_status["start_time"]
+        hours, remainder = divmod(int(delta.total_seconds()), 3600)
+        minutes, seconds = divmod(remainder, 60)
+        elapsed = f"{hours}h {minutes}m {seconds}s"
+    
+    # Pre-compute formatted loss string to avoid f-string format specifier issues
+    if training_status["total_loss"]:
+        loss_str = f"{training_status['total_loss']:.4f}"
+    else:
+        loss_str = "N/A"
+    
+    # Pre-compute other conditional strings
+    state_str = "🟢 Training" if training_status["is_training"] else "🔴 Stopped"
+    continuous_str = "✅ Enabled" if CONTINUOUS_TRAINING else "❌ Disabled"
+    elapsed_str = elapsed if elapsed else "N/A"
+    effective_batch = BATCH_SIZE * GRADIENT_ACCUMULATION
+    
+    return f"""
+## 🤖 AI Python Model Trainer
+
+### Status
+| Item | Value |
+|------|-------|
+| **State** | {state_str} |
+| **Message** | {training_status["message"]} |
+| **Sessions Completed** | {training_status["session_count"]} |
+
+### Progress
+| Metric | Value |
+|--------|-------|
+| **Current Step** | {training_status["current_step"]:,} / {MAX_STEPS_PER_SESSION:,} |
+| **Current Loss** | {loss_str} |
+| **Last Checkpoint** | {training_status["last_save"]} |
+| **Elapsed Time** | {elapsed_str} |
+
+### Configuration
+| Setting | Value |
+|---------|-------|
+| **Model Repo** | [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO}) |
+| **Continuous Mode** | {continuous_str} |
+| **Batch Size** | {BATCH_SIZE} (effective: {effective_batch}) |
+| **Max Steps/Session** | {MAX_STEPS_PER_SESSION:,} |
+"""
+
+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)
+    thread.start()
+    return "🚀 Training started in background!"
+
+def stop_training():
+    """Stop training"""
+    global stop_requested
+    if not training_status["is_training"]:
+        return "⚠️ No training in progress"
+    stop_requested = True
+    training_status["message"] = "⏹️ Stopping after current step..."
+    return "⏹️ Stop requested"
+
+# ============ 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:
+            print("[INFO] Auto-starting training session...")
+            start_training()
+        time.sleep(WAIT_BETWEEN_SESSIONS)
+
+auto_thread = threading.Thread(target=auto_start, daemon=True)
+auto_thread.start()
+
+# ============ GRADIO APP ============
+with gr.Blocks(title="AI Python Trainer") as demo:
+    gr.Markdown("# 🐍 AI Python Code Model Trainer")
+    gr.Markdown(f"**Continuous training** on `{DATASET_NAME}` with auto-checkpoint")
+    
+    status_display = gr.Markdown(get_status)
+    
+    with gr.Row():
+        start_btn = gr.Button("▶️ Start Training", variant="primary")
+        stop_btn = gr.Button("⏹️ Stop Training", variant="stop")
+        refresh_btn = gr.Button("🔄 Refresh Status")
+    
+    output = gr.Textbox(label="Output", interactive=False)
+    
+    start_btn.click(start_training_async, outputs=output)
+    stop_btn.click(stop_training, outputs=output)
+    refresh_btn.click(get_status, outputs=status_display)
+    
+    demo.load(get_status, outputs=status_display, every=30)
+
+demo.launch()