Upload 2 files

app.py

@@ -0,0 +1,380 @@
+"""
+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()

requirements.txt

@@ -0,0 +1,6 @@
+transformers>=4.36.0
+datasets>=2.16.0
+accelerate>=0.25.0
+gradio>=4.0.0
+huggingface_hub>=0.20.0
+torch>=2.0.0