test

app.py

@@ -376,7 +376,311 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
                 shutil.rmtree(temp_dir)
         except:
             pass
+def train_model_background(job_id, dataset_name, base_model_name=None):
+    """Background training function with improved configuration"""
+    progress = training_jobs[job_id]
+    
+    try:
+        # Create a temporary directory for this job
+        temp_dir = tempfile.mkdtemp(prefix=f"train_{job_id}_")
+        
+        # Set environment variables for caching
+        os.environ['HF_HOME'] = temp_dir
+        os.environ['TRANSFORMERS_CACHE'] = temp_dir
+        os.environ['HF_DATASETS_CACHE'] = temp_dir
+        os.environ['TORCH_HOME'] = temp_dir
+        
+        progress.status = "loading_libraries"
+        progress.message = "Loading required libraries..."
+        
+        # Import heavy libraries after setting cache paths
+        import torch
+        from datasets import load_dataset, Dataset
+        from huggingface_hub import login
+        from transformers import (
+            AutoModelForCausalLM,
+            AutoTokenizer,
+            TrainingArguments,
+            Trainer,
+            TrainerCallback,
+            DataCollatorForLanguageModeling
+        )
+        from peft import (
+            LoraConfig,
+            get_peft_model,
+            TaskType
+        )
+        
+        # === Authentication ===
+        hf_token = os.getenv('HF_TOKEN')
+        if hf_token:
+            login(token=hf_token)
+        
+        progress.status = "loading_model"
+        progress.message = "Loading base model and tokenizer..."
+
+        # === Better Model Selection ===
+        # Use a more suitable model for medical conversations
+        base_model = base_model_name or "microsoft/DialoGPT-medium"  # Better than small
+        new_model = f"trained-model-{job_id}"
+        max_length = 512  # Increased for better context
+        
+        # === Load Model and Tokenizer ===
+        model = AutoModelForCausalLM.from_pretrained(
+            base_model,
+            cache_dir=temp_dir,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else "cpu",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            base_model,
+            cache_dir=temp_dir,
+            trust_remote_code=True,
+            padding_side="right"  # Important for causal LM
+        )
+        
+        # Add padding token if not present
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+        
+        # Resize token embeddings if needed
+        model.resize_token_embeddings(len(tokenizer))
+
+        progress.status = "preparing_model"
+        progress.message = "Setting up improved LoRA configuration..."
+
+        # === Better LoRA Config ===
+        peft_config = LoraConfig(
+            r=16,  # Increased rank for better learning
+            lora_alpha=32,  # Increased alpha
+            lora_dropout=0.05,  # Reduced dropout
+            bias="none",
+            task_type=TaskType.CAUSAL_LM,
+            target_modules=["c_attn", "c_proj"],  # Target specific modules for DialoGPT
+        )
+        model = get_peft_model(model, peft_config)
+        
+        # Print trainable parameters
+        model.print_trainable_parameters()
+
+        progress.status = "loading_dataset"
+        progress.message = "Loading and preparing dataset..."
+
+        # === Load & Prepare Dataset ===
+        dataset = load_dataset(
+            dataset_name, 
+            split="train" if "train" in load_dataset(dataset_name, cache_dir=temp_dir).keys() else "all",
+            cache_dir=temp_dir,
+            trust_remote_code=True
+        )
+        
+        # Automatically detect question and answer columns
+        question_col, answer_col = detect_qa_columns(dataset)
+        
+        if not question_col or not answer_col:
+            raise ValueError("Could not automatically detect question and answer columns in the dataset")
+        
+        progress.detected_columns = {"question": question_col, "answer": answer_col}
+        progress.message = f"Detected columns - Question: {question_col}, Answer: {answer_col}"
+        
+        # Use more data for better training
+        dataset_size = min(1000, len(dataset))  # Increased from 100 to 1000
+        dataset = dataset.shuffle(seed=42).select(range(dataset_size))
+
+        # === Better Text Formatting ===
+        def format_conversation(example):
+            question = str(example[question_col]).strip()
+            answer = str(example[answer_col]).strip()
+            
+            # Better formatting with special tokens
+            conversation = f"<|user|>{question}<|assistant|>{answer}<|endoftext|>"
+            return {"text": conversation}
+        
+        # Apply formatting
+        dataset = dataset.map(format_conversation, remove_columns=dataset.column_names)
+        
+        # Filter out very short or very long examples
+        dataset = dataset.filter(lambda x: 10 < len(x["text"]) < max_length * 2)
+
+        # === Improved Training Arguments ===
+        batch_size = 4 if torch.cuda.is_available() else 2
+        gradient_accumulation_steps = 2
+        num_epochs = 3  # Increased epochs
+        learning_rate = 2e-4  # Better learning rate
+        
+        steps_per_epoch = len(dataset) // (batch_size * gradient_accumulation_steps)
+        total_steps = steps_per_epoch * num_epochs
+        warmup_steps = max(10, total_steps // 10)  # 10% warmup
+        
+        progress.total_steps = total_steps
+        progress.status = "training"
+        progress.message = "Starting improved training..."
+
+        output_dir = os.path.join(temp_dir, new_model)
+        os.makedirs(output_dir, exist_ok=True)
+        
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            num_train_epochs=num_epochs,
+            learning_rate=learning_rate,
+            warmup_steps=warmup_steps,
+            logging_steps=5,
+            save_steps=max(10, total_steps // 4),
+            save_total_limit=2,
+            evaluation_strategy="no",
+            logging_strategy="steps",
+            save_strategy="steps",
+            fp16=torch.cuda.is_available(),
+            bf16=False,
+            dataloader_num_workers=0,
+            remove_unused_columns=False,
+            report_to=None,
+            prediction_loss_only=True,
+            optim="adamw_torch",
+            weight_decay=0.01,
+            lr_scheduler_type="cosine",
+            gradient_checkpointing=True,
+            dataloader_pin_memory=False,
+        )
+
+        # === Data Collator ===
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False,  # We're doing causal LM, not masked LM
+            return_tensors="pt",
+            pad_to_multiple_of=8,
+        )
+
+        # Custom tokenization function
+        def tokenize_function(examples):
+            # Tokenize the text
+            tokenized = tokenizer(
+                examples["text"],
+                truncation=True,
+                padding=False,  # Will be handled by data collator
+                max_length=max_length,
+                return_tensors=None,
+            )
+            
+            # For causal LM, labels are the same as input_ids
+            tokenized["labels"] = tokenized["input_ids"].copy()
+            return tokenized
+
+        # Tokenize dataset
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=dataset.column_names,
+            desc="Tokenizing dataset",
+        )
+
+        # Custom callback to track progress
+        class ProgressCallback(TrainerCallback):
+            def __init__(self, progress_tracker):
+                self.progress_tracker = progress_tracker
+                self.last_update = time.time()
+            
+            def on_log(self, args, state, control, model=None, logs=None, **kwargs):
+                current_time = time.time()
+                # Update every 5 seconds
+                if current_time - self.last_update >= 5:
+                    self.progress_tracker.update_progress(
+                        state.global_step, 
+                        state.max_steps,
+                        f"Training step {state.global_step}/{state.max_steps}"
+                    )
+                    self.last_update = current_time
+                    
+                    # Log training metrics if available
+                    if logs:
+                        loss = logs.get('train_loss', logs.get('loss', 'N/A'))
+                        lr = logs.get('learning_rate', 'N/A')
+                        self.progress_tracker.message = f"Step {state.global_step}/{state.max_steps}, Loss: {loss:.4f}, LR: {lr}"
+            
+            def on_train_begin(self, args, state, control, **kwargs):
+                self.progress_tracker.status = "training"
+                self.progress_tracker.message = "Training started with improved configuration..."
+            
+            def on_train_end(self, args, state, control, **kwargs):
+                self.progress_tracker.status = "saving"
+                self.progress_tracker.message = "Training complete, saving improved model..."
 
+        # === Trainer Initialization ===
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset,
+            data_collator=data_collator,
+            callbacks=[ProgressCallback(progress)],
+            tokenizer=tokenizer,
+        )
+
+        # === Train & Save ===
+        trainer.train()
+        
+        # Save the model properly
+        trainer.save_model(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        
+        # Also save the base model name for inference
+        with open(os.path.join(output_dir, "base_model.txt"), "w") as f:
+            f.write(base_model)
+        
+        # Save training info
+        training_info = {
+            "base_model": base_model,
+            "dataset_name": dataset_name,
+            "dataset_size": len(dataset),
+            "max_length": max_length,
+            "batch_size": batch_size,
+            "learning_rate": learning_rate,
+            "num_epochs": num_epochs,
+            "total_steps": total_steps,
+            "detected_columns": progress.detected_columns
+        }
+        
+        with open(os.path.join(output_dir, "training_info.json"), "w") as f:
+            import json
+            json.dump(training_info, f, indent=2)
+        
+        # Save model info
+        progress.model_path = output_dir
+        progress.status = "completed"
+        progress.progress = 100
+        progress.message = f"Improved training completed! Model ready for download."
+        
+        # Keep the temp directory for download (cleanup after 2 hours for larger model)
+        def cleanup_temp_dir():
+            time.sleep(7200)  # Wait 2 hours before cleanup
+            try:
+                shutil.rmtree(temp_dir)
+                # Remove from training_jobs after cleanup
+                if job_id in training_jobs:
+                    del training_jobs[job_id]
+            except:
+                pass
+        
+        cleanup_thread = threading.Thread(target=cleanup_temp_dir)
+        cleanup_thread.daemon = True
+        cleanup_thread.start()
+        
+    except Exception as e:
+        progress.status = "error"
+        progress.error = str(e)
+        progress.message = f"Training failed: {str(e)}"
+        
+        # Clean up on error
+        try:
+            if 'temp_dir' in locals():
+                shutil.rmtree(temp_dir)
+        except:
+            pass
 def create_model_zip(model_path, job_id):
     """Create a zip file containing the trained model"""
     memory_file = io.BytesIO()