Update app.py

app.py

@@ -222,57 +222,87 @@ def log_message(output_log, msg):
 # 🧠 Train model to expand short prompts into long ones
 # =====================================================
 @spaces.GPU(duration=300)
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+)
+from peft import LoraConfig, get_peft_model, TaskType
+import os
+
+# =====================================================
+# 🔧 Utility logging
+# =====================================================
+def log_message(output_log, msg):
+    print(msg)
+    output_log.append(msg)
+
+def start_async_upload(output_dir, hf_repo, output_log):
+    from huggingface_hub import upload_folder
+    try:
+        upload_folder(
+            repo_id=hf_repo,
+            folder_path=output_dir,
+            repo_type="model",
+            commit_message="Upload fine-tuned model"
+        )
+        log_message(output_log, f"☁️ Model uploaded to {hf_repo}")
+    except Exception as e:
+        log_message(output_log, f"⚠️ Upload failed: {e}")
+
+# =====================================================
+# 🧠 GPU-safe training for short→long prompt expansion
+# =====================================================
 def train_model(
     base_model: str,
     dataset_name: str,
-    num_epochs: int,
-    batch_size: int,
-    learning_rate: float,
-    hf_repo: str,
+    num_epochs: int = 1,
+    batch_size: int = 1,
+    learning_rate: float = 2e-4,
+    hf_repo: str = None,
 ):
     output_log = []
 
     try:
         log_message(output_log, "🚀 Initializing prompt expansion training...")
 
-        # ===== Device =====
+        # ===== Device setup =====
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        log_message(output_log, f"🎮 Using device: {device}")
+        dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+        log_message(output_log, f"🎮 Device: {device}, dtype: {dtype}")
+
         if device == "cuda":
-            log_message(output_log, f"✅ GPU: {torch.cuda.get_device_name(0)}")
+            gpu_name = torch.cuda.get_device_name(0)
+            gpu_mem = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+            log_message(output_log, f"✅ GPU: {gpu_name} ({gpu_mem:.1f} GB)")
 
-        # ===== Load dataset =====
+        # ===== Load dataset safely =====
         log_message(output_log, f"\n📚 Loading dataset: {dataset_name}")
         dataset = load_dataset(dataset_name)
-        dataset = dataset["train"].train_test_split(test_size=0.2)
+        dataset = dataset["train"].train_test_split(test_size=0.2, seed=42)
         train_dataset = dataset["train"]
         test_dataset = dataset["test"]
+        log_message(output_log, f"   → Train samples: {len(train_dataset)} | Test samples: {len(test_dataset)}")
 
-        log_message(output_log, f"   → Train samples: {len(train_dataset)}")
-        log_message(output_log, f"   → Test samples:  {len(test_dataset)}")
-        log_message(output_log, f"   → Columns: {train_dataset.column_names}")
-
-        # =====================================================
-        # 🧩 Format training examples
-        # Each sample has 'short' (input) and 'long' (target)
-        # =====================================================
+        # ===== Format examples =====
         def format_example(example):
-            short_prompt = example.get("short_prompt", "").strip()
-            long_response = example.get("long_prompt", "").strip()
+            short_prompt = example.get("short", "").strip()
+            long_response = example.get("long", "").strip()
 
-            # Compose a structured conversation
             prompt = (
-                f"<|system|>\nYou are a helpful AI that expands short prompts into rich, detailed prompts.\n"
+                f"<|system|>\nYou are an AI that expands short prompts into detailed, descriptive versions.\n"
                 f"<|user|>\nShort: {short_prompt}\n"
                 f"<|assistant|>\n{long_response}"
             )
             return {"text": prompt}
 
-        train_dataset = train_dataset.map(format_example)
-        test_dataset = test_dataset.map(format_example)
-        log_message(output_log, f"✅ Prepared {len(train_dataset)} train + {len(test_dataset)} test samples")
+        train_dataset = train_dataset.map(format_example, num_proc=1)
+        test_dataset = test_dataset.map(format_example, num_proc=1)
 
-        # ===== Load model & tokenizer =====
+        # ===== Load model & tokenizer safely =====
         log_message(output_log, f"\n🤖 Loading base model: {base_model}")
         tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
         if tokenizer.pad_token is None:
@@ -280,55 +310,67 @@ def train_model(
 
         model = AutoModelForCausalLM.from_pretrained(
             base_model,
+            torch_dtype=dtype,
             trust_remote_code=True,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
             low_cpu_mem_usage=True,
-        ).to(device)
+            device_map="auto" if device == "cuda" else None,
+        )
 
-        log_message(output_log, "✅ Model and tokenizer loaded successfully")
+        # Enable memory optimizations
+        model.gradient_checkpointing_enable()
+        log_message(output_log, "✅ Model loaded with gradient checkpointing")
 
-        # ===== LoRA fine-tuning setup =====
-        log_message(output_log, "\n⚙️ Applying LoRA configuration...")
+        # ===== Apply lightweight LoRA =====
+        log_message(output_log, "\n⚙️ Applying LoRA fine-tuning config...")
         lora_config = LoraConfig(
             task_type=TaskType.CAUSAL_LM,
-            r=8,
-            lora_alpha=16,
+            r=4,
+            lora_alpha=8,
             lora_dropout=0.1,
             target_modules=["q_proj", "v_proj"],
             bias="none",
         )
         model = get_peft_model(model, lora_config)
-        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        log_message(output_log, f"Trainable parameters: {trainable_params:,}")
+        log_message(output_log, f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
 
-        # ===== Tokenization =====
+        # ===== Tokenization (batched for speed) =====
         def tokenize_fn(examples):
             tokenized = tokenizer(
                 examples["text"],
                 padding="max_length",
                 truncation=True,
-                max_length=512,
+                max_length=384,
             )
             tokenized["labels"] = tokenized["input_ids"].copy()
             return tokenized
 
         train_dataset = train_dataset.map(tokenize_fn, batched=True)
         test_dataset = test_dataset.map(tokenize_fn, batched=True)
-        log_message(output_log, "✅ Tokenization complete")
+        log_message(output_log, "✅ Tokenization done")
 
         # ===== Training setup =====
         output_dir = "./prompt_expander_lora"
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Automatically reduce batch size for low GPU VRAM
+        if device == "cuda" and gpu_mem < 10:
+            batch_size = 1
+            log_message(output_log, f"⚠️ GPU memory low → Using batch_size={batch_size}")
+
         training_args = TrainingArguments(
             output_dir=output_dir,
             num_train_epochs=num_epochs,
             per_device_train_batch_size=batch_size,
-            gradient_accumulation_steps=2,
-            warmup_steps=10,
-            logging_steps=5,
+            gradient_accumulation_steps=4,
+            warmup_steps=20,
+            logging_steps=10,
             save_strategy="epoch",
-            fp16=device == "cuda",
             optim="adamw_torch",
             learning_rate=learning_rate,
+            fp16=(dtype == torch.float16),
+            bf16=(dtype == torch.bfloat16),
+            max_grad_norm=1.0,
+            report_to="none",
         )
 
         trainer = Trainer(
@@ -340,22 +382,25 @@ def train_model(
         )
 
         # ===== Train =====
-        log_message(output_log, "\n🔥 Starting training...")
+        log_message(output_log, "\n🔥 Starting safe LoRA fine-tuning...")
         trainer.train()
 
-        # ===== Save model =====
-        log_message(output_log, "\n💾 Saving fine-tuned model locally...")
+        # ===== Save =====
+        log_message(output_log, "\n💾 Saving fine-tuned model...")
         trainer.save_model(output_dir)
         tokenizer.save_pretrained(output_dir)
 
         # ===== Upload to Hub =====
-        log_message(output_log, f"\n☁️ Uploading to Hugging Face Hub: {hf_repo}")
-        start_async_upload(output_dir, hf_repo, output_log)
+        if hf_repo:
+            log_message(output_log, f"\n☁️ Uploading to {hf_repo} ...")
+            start_async_upload(output_dir, hf_repo, output_log)
 
-        log_message(output_log, "\n✅ Training complete and model uploaded!")
+        log_message(output_log, "\n✅ Training complete!")
 
+    except torch.cuda.OutOfMemoryError:
+        log_message(output_log, "\n❌ CUDA OOM — try lowering batch size or sequence length.")
     except Exception as e:
-        log_message(output_log, f"\n❌ Error during training: {e}")
+        log_message(output_log, f"\n❌ Training error: {e}")
 
     return "\n".join(output_log)