test

app.py

@@ -112,272 +112,7 @@ def detect_qa_columns(dataset):
     return question_col, answer_col
 
 def train_model_background(job_id, dataset_name, base_model_name=None):
-    """Background training function with progress tracking"""
-    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,
-        )
-        from peft import (
-            LoraConfig,
-            get_peft_model,
-        )
-        
-        # === 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..."
-
-        # === Configuration ===
-        base_model = base_model_name or "microsoft/DialoGPT-small"
-        new_model = f"trained-model-{job_id}"
-        max_length = 256
-        
-        # === Load Model and Tokenizer ===
-        model = AutoModelForCausalLM.from_pretrained(
-            base_model,
-            cache_dir=temp_dir,
-            torch_dtype=torch.float32,
-            device_map="auto" if torch.cuda.is_available() else "cpu",
-            trust_remote_code=True
-        )
-        
-        tokenizer = AutoTokenizer.from_pretrained(
-            base_model,
-            cache_dir=temp_dir,
-            trust_remote_code=True
-        )
-        
-        # Add padding token if not present
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-        
-        # Resize token embeddings if needed
-        model.resize_token_embeddings(len(tokenizer))
-
-        progress.status = "preparing_model"
-        progress.message = "Setting up LoRA configuration..."
-
-        # === LoRA Config ===
-        peft_config = LoraConfig(
-            r=8,
-            lora_alpha=16,
-            lora_dropout=0.1,
-            bias="none",
-            task_type="CAUSAL_LM",
-        )
-        model = get_peft_model(model, peft_config)
-
-        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 subset for faster testing (can be made configurable)
-        dataset = dataset.shuffle(seed=65).select(range(min(100, len(dataset))))
-
-        # Custom dataset class for proper handling
-        class CustomDataset(torch.utils.data.Dataset):
-            def __init__(self, texts, tokenizer, max_length):
-                self.texts = texts
-                self.tokenizer = tokenizer
-                self.max_length = max_length
-
-            def __len__(self):
-                return len(self.texts)
-
-            def __getitem__(self, idx):
-                text = self.texts[idx]
-                
-                # Tokenize the text
-                encoding = self.tokenizer(
-                    text,
-                    truncation=True,
-                    padding='max_length',
-                    max_length=self.max_length,
-                    return_tensors='pt'
-                )
-                
-                # Flatten the tensors (remove batch dimension)
-                input_ids = encoding['input_ids'].squeeze()
-                attention_mask = encoding['attention_mask'].squeeze()
-                
-                # For causal language modeling, labels are the same as input_ids
-                labels = input_ids.clone()
-                
-                # Set labels to -100 for padding tokens (they won't contribute to loss)
-                labels[attention_mask == 0] = -100
-                
-                return {
-                    'input_ids': input_ids,
-                    'attention_mask': attention_mask,
-                    'labels': labels
-                }
-
-        # Prepare texts using detected columns
-        texts = []
-        for item in dataset:
-            question = str(item[question_col]).strip()
-            answer = str(item[answer_col]).strip()
-            text = f"Question: {question}\nAnswer: {answer}{tokenizer.eos_token}"
-            texts.append(text)
-
-        # Create custom dataset
-        train_dataset = CustomDataset(texts, tokenizer, max_length)
-
-        # Calculate total training steps
-        batch_size = 2
-        gradient_accumulation_steps = 1
-        num_epochs = 1
-        
-        steps_per_epoch = len(train_dataset) // (batch_size * gradient_accumulation_steps)
-        total_steps = steps_per_epoch * num_epochs
-        
-        progress.total_steps = total_steps
-        progress.status = "training"
-        progress.message = "Starting training..."
-
-        # === Training Arguments ===
-        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,
-            logging_steps=1,
-            save_steps=max(1, total_steps // 2),
-            save_total_limit=1,
-            learning_rate=5e-5,
-            warmup_steps=2,
-            logging_strategy="steps",
-            save_strategy="steps",
-            fp16=False,
-            bf16=False,
-            dataloader_num_workers=0,
-            remove_unused_columns=False,
-            report_to=None,
-            prediction_loss_only=True,
-        )
-
-        # 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 3 seconds
-                if current_time - self.last_update >= 3:
-                    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'))
-                        self.progress_tracker.message = f"Step {state.global_step}/{state.max_steps}, Loss: {loss}"
-            
-            def on_train_begin(self, args, state, control, **kwargs):
-                self.progress_tracker.status = "training"
-                self.progress_tracker.message = "Training started..."
-            
-            def on_train_end(self, args, state, control, **kwargs):
-                self.progress_tracker.status = "saving"
-                self.progress_tracker.message = "Training complete, saving model..."
-
-        # === Trainer Initialization ===
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_dataset,
-            callbacks=[ProgressCallback(progress)],
-            tokenizer=tokenizer,
-        )
-
-        # === Train & Save ===
-        trainer.train()
-        trainer.save_model(output_dir)
-        tokenizer.save_pretrained(output_dir)
-        
-        # Save model info
-        progress.model_path = output_dir
-        progress.status = "completed"
-        progress.progress = 100
-        progress.message = f"Training completed! Model ready for download."
-        
-        # Keep the temp directory for download (cleanup after 1 hour)
-        def cleanup_temp_dir():
-            time.sleep(3600)  # Wait 1 hour 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 train_model_background(job_id, dataset_name, base_model_name=None):
-    """Background training function with improved configuration"""
+    """Background training function with fixed tokenization"""
     progress = training_jobs[job_id]
     
     try:
@@ -419,11 +154,10 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
         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
+        # === Model Configuration ===
+        base_model = base_model_name or "microsoft/DialoGPT-medium"
         new_model = f"trained-model-{job_id}"
-        max_length = 512  # Increased for better context
+        max_length = 512
         
         # === Load Model and Tokenizer ===
         model = AutoModelForCausalLM.from_pretrained(
@@ -439,7 +173,7 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
             base_model,
             cache_dir=temp_dir,
             trust_remote_code=True,
-            padding_side="right"  # Important for causal LM
+            padding_side="right"
         )
         
         # Add padding token if not present
@@ -451,16 +185,16 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
         model.resize_token_embeddings(len(tokenizer))
 
         progress.status = "preparing_model"
-        progress.message = "Setting up improved LoRA configuration..."
+        progress.message = "Setting up LoRA configuration..."
 
-        # === Better LoRA Config ===
+        # === LoRA Config ===
         peft_config = LoraConfig(
-            r=16,  # Increased rank for better learning
-            lora_alpha=32,  # Increased alpha
-            lora_dropout=0.05,  # Reduced dropout
+            r=16,
+            lora_alpha=32,
+            lora_dropout=0.05,
             bias="none",
             task_type=TaskType.CAUSAL_LM,
-            target_modules=["c_attn", "c_proj"],  # Target specific modules for DialoGPT
+            target_modules=["c_attn", "c_proj"],
         )
         model = get_peft_model(model, peft_config)
         
@@ -487,38 +221,61 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
         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
+        # Use subset for faster training
+        dataset_size = min(500, len(dataset))
         dataset = dataset.shuffle(seed=42).select(range(dataset_size))
 
-        # === Better Text Formatting ===
+        # === Fixed 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|>"
+            # Simple format that works well with tokenizer
+            conversation = f"Question: {question}\nAnswer: {answer}{tokenizer.eos_token}"
             return {"text": conversation}
         
         # Apply formatting
-        dataset = dataset.map(format_conversation, remove_columns=dataset.column_names)
+        formatted_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)
+        formatted_dataset = formatted_dataset.filter(lambda x: 10 < len(x["text"]) < max_length * 3)
+
+        # === Fixed Tokenization Function ===
+        def tokenize_function(examples):
+            # Tokenize the text
+            model_inputs = 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
+            model_inputs["labels"] = model_inputs["input_ids"].copy()
+            return model_inputs
 
-        # === Improved Training Arguments ===
+        # Tokenize dataset
+        tokenized_dataset = formatted_dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=formatted_dataset.column_names,
+            desc="Tokenizing dataset",
+        )
+
+        # === Training Configuration ===
         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
+        num_epochs = 2
+        learning_rate = 2e-4
         
-        steps_per_epoch = len(dataset) // (batch_size * gradient_accumulation_steps)
+        steps_per_epoch = len(tokenized_dataset) // (batch_size * gradient_accumulation_steps)
         total_steps = steps_per_epoch * num_epochs
-        warmup_steps = max(10, total_steps // 10)  # 10% warmup
+        warmup_steps = max(10, total_steps // 10)
         
         progress.total_steps = total_steps
         progress.status = "training"
-        progress.message = "Starting improved training..."
+        progress.message = "Starting training..."
 
         output_dir = os.path.join(temp_dir, new_model)
         os.makedirs(output_dir, exist_ok=True)
@@ -552,32 +309,9 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
         # === Data Collator ===
         data_collator = DataCollatorForLanguageModeling(
             tokenizer=tokenizer,
-            mlm=False,  # We're doing causal LM, not masked LM
+            mlm=False,
             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",
+            pad_to_multiple_of=8 if torch.cuda.is_available() else None,
         )
 
         # Custom callback to track progress
@@ -588,7 +322,6 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
             
             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, 
@@ -597,19 +330,20 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
                     )
                     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}"
+                        if isinstance(loss, (int, float)):
+                            loss = f"{loss:.4f}"
+                        self.progress_tracker.message = f"Step {state.global_step}/{state.max_steps}, Loss: {loss}, 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..."
+                self.progress_tracker.message = "Training started..."
             
             def on_train_end(self, args, state, control, **kwargs):
                 self.progress_tracker.status = "saving"
-                self.progress_tracker.message = "Training complete, saving improved model..."
+                self.progress_tracker.message = "Training complete, saving model..."
 
         # === Trainer Initialization ===
         trainer = Trainer(
@@ -628,15 +362,14 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
         trainer.save_model(output_dir)
         tokenizer.save_pretrained(output_dir)
         
-        # Also save the base model name for inference
+        # Save additional info
         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),
+            "dataset_size": len(tokenized_dataset),
             "max_length": max_length,
             "batch_size": batch_size,
             "learning_rate": learning_rate,
@@ -649,18 +382,17 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
             import json
             json.dump(training_info, f, indent=2)
         
-        # Save model info
+        # Update progress
         progress.model_path = output_dir
         progress.status = "completed"
         progress.progress = 100
-        progress.message = f"Improved training completed! Model ready for download."
+        progress.message = f"Training completed successfully! Model ready for download."
         
-        # Keep the temp directory for download (cleanup after 2 hours for larger model)
+        # Keep the temp directory for download
         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:
@@ -681,6 +413,7 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
                 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()
@@ -694,6 +427,7 @@ def create_model_zip(model_path, job_id):
     
     memory_file.seek(0)
     return memory_file
+
 # ============== API ROUTES ==============
 @app.route('/api/train', methods=['POST'])
 def start_training():
@@ -701,9 +435,9 @@ def start_training():
     try:
         data = request.get_json() if request.is_json else {}
         dataset_name = data.get('dataset_name', 'ruslanmv/ai-medical-chatbot')
-        base_model_name = data.get('base_model', 'microsoft/DialoGPT-small')
+        base_model_name = data.get('base_model', 'microsoft/DialoGPT-medium')
         
-        job_id = str(uuid.uuid4())[:8]  # Short UUID
+        job_id = str(uuid.uuid4())[:8]
         progress = TrainingProgress(job_id)
         training_jobs[job_id] = progress
         
@@ -797,7 +531,7 @@ def home():
                 "url": "/api/train",
                 "body": {
                     "dataset_name": "your-dataset-name",
-                    "base_model": "microsoft/DialoGPT-small"
+                    "base_model": "microsoft/DialoGPT-medium"
                 }
             }
         }
@@ -808,5 +542,5 @@ def health():
     return jsonify({"status": "healthy"})
 
 if __name__ == '__main__':
-    port = int(os.environ.get('PORT', 7860))  # HF Spaces uses port 7860
+    port = int(os.environ.get('PORT', 7860))
     app.run(host='0.0.0.0', port=port, debug=False)