test

app.py

@@ -112,7 +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 fixed tokenization"""
+    """Background training function with progress tracking"""
     progress = training_jobs[job_id]
     
     try:
@@ -138,12 +138,10 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
             TrainingArguments,
             Trainer,
             TrainerCallback,
-            DataCollatorForLanguageModeling
         )
         from peft import (
             LoraConfig,
             get_peft_model,
-            TaskType
         )
         
         # === Authentication ===
@@ -154,32 +152,29 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
         progress.status = "loading_model"
         progress.message = "Loading base model and tokenizer..."
 
-        # === Model Configuration ===
-        base_model = base_model_name or "microsoft/DialoGPT-medium"
+        # === Configuration ===
+        base_model = base_model_name or "microsoft/DialoGPT-small"
         new_model = f"trained-model-{job_id}"
-        max_length = 512
+        max_length = 256
         
         # === 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,
+            torch_dtype=torch.float32,
             device_map="auto" if torch.cuda.is_available() else "cpu",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True
+            trust_remote_code=True
         )
         
         tokenizer = AutoTokenizer.from_pretrained(
             base_model,
             cache_dir=temp_dir,
-            trust_remote_code=True,
-            padding_side="right"
+            trust_remote_code=True
         )
         
         # 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))
@@ -189,17 +184,13 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
 
         # === LoRA Config ===
         peft_config = LoraConfig(
-            r=16,
-            lora_alpha=32,
-            lora_dropout=0.05,
+            r=8,
+            lora_alpha=16,
+            lora_dropout=0.1,
             bias="none",
-            task_type=TaskType.CAUSAL_LM,
-            target_modules=["c_attn", "c_proj"],
+            task_type="CAUSAL_LM",
         )
         model = get_peft_model(model, peft_config)
-        
-        # Print trainable parameters
-        model.print_trainable_parameters()
 
         progress.status = "loading_dataset"
         progress.message = "Loading and preparing dataset..."
@@ -221,62 +212,71 @@ 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 subset for faster training
-        dataset_size = min(500, len(dataset))
-        dataset = dataset.shuffle(seed=42).select(range(dataset_size))
+        # Use subset for faster testing (can be made configurable)
+        dataset = dataset.shuffle(seed=65).select(range(min(1000, len(dataset))))
 
-        # === Fixed Text Formatting ===
-        def format_conversation(example):
-            question = str(example[question_col]).strip()
-            answer = str(example[answer_col]).strip()
-            
-            # Simple format that works well with tokenizer
-            conversation = f"Question: {question}\nAnswer: {answer}{tokenizer.eos_token}"
-            return {"text": conversation}
-        
-        # Apply formatting
-        formatted_dataset = dataset.map(format_conversation, remove_columns=dataset.column_names)
-        
-        # Filter out very short or very long examples
-        formatted_dataset = formatted_dataset.filter(lambda x: 10 < len(x["text"]) < max_length * 3)
+        # 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
 
-        # === 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
+            def __len__(self):
+                return len(self.texts)
 
-        # Tokenize dataset
-        tokenized_dataset = formatted_dataset.map(
-            tokenize_function,
-            batched=True,
-            remove_columns=formatted_dataset.column_names,
-            desc="Tokenizing dataset",
-        )
+            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)
 
-        # === Training Configuration ===
-        batch_size = 4 if torch.cuda.is_available() else 2
-        gradient_accumulation_steps = 2
-        num_epochs = 2
-        learning_rate = 2e-4
+        # Calculate total training steps
+        batch_size = 2
+        gradient_accumulation_steps = 1
+        num_epochs = 1
         
-        steps_per_epoch = len(tokenized_dataset) // (batch_size * gradient_accumulation_steps)
+        steps_per_epoch = len(train_dataset) // (batch_size * gradient_accumulation_steps)
         total_steps = steps_per_epoch * num_epochs
-        warmup_steps = max(10, total_steps // 10)
         
         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)
         
@@ -285,33 +285,19 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
             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_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=torch.cuda.is_available(),
+            fp16=False,
             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,
-            return_tensors="pt",
-            pad_to_multiple_of=8 if torch.cuda.is_available() else None,
         )
 
         # Custom callback to track progress
@@ -322,7 +308,8 @@ 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()
-                if current_time - self.last_update >= 5:
+                # Update every 3 seconds
+                if current_time - self.last_update >= 3:
                     self.progress_tracker.update_progress(
                         state.global_step, 
                         state.max_steps,
@@ -330,12 +317,10 @@ 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')
-                        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}"
+                        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"
@@ -349,50 +334,28 @@ def train_model_background(job_id, dataset_name, base_model_name=None):
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=tokenized_dataset,
-            data_collator=data_collator,
+            train_dataset=train_dataset,
             callbacks=[ProgressCallback(progress)],
             tokenizer=tokenizer,
         )
 
         # === Train & Save ===
         trainer.train()
-        
-        # Save the model properly
         trainer.save_model(output_dir)
         tokenizer.save_pretrained(output_dir)
         
-        # Save additional info
-        with open(os.path.join(output_dir, "base_model.txt"), "w") as f:
-            f.write(base_model)
-        
-        training_info = {
-            "base_model": base_model,
-            "dataset_name": dataset_name,
-            "dataset_size": len(tokenized_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)
-        
-        # Update progress
+        # Save model info
         progress.model_path = output_dir
         progress.status = "completed"
         progress.progress = 100
-        progress.message = f"Training completed successfully! Model ready for download."
+        progress.message = f"Training completed! Model ready for download."
         
-        # Keep the temp directory for download
+        # Keep the temp directory for download (cleanup after 1 hour)
         def cleanup_temp_dir():
-            time.sleep(7200)  # Wait 2 hours before cleanup
+            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:
@@ -427,7 +390,6 @@ 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():
@@ -435,9 +397,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-medium')
+        base_model_name = data.get('base_model', 'microsoft/DialoGPT-small')
         
-        job_id = str(uuid.uuid4())[:8]
+        job_id = str(uuid.uuid4())[:8]  # Short UUID
         progress = TrainingProgress(job_id)
         training_jobs[job_id] = progress
         
@@ -531,7 +493,7 @@ def home():
                 "url": "/api/train",
                 "body": {
                     "dataset_name": "your-dataset-name",
-                    "base_model": "microsoft/DialoGPT-medium"
+                    "base_model": "microsoft/DialoGPT-small"
                 }
             }
         }
@@ -542,5 +504,5 @@ def health():
     return jsonify({"status": "healthy"})
 
 if __name__ == '__main__':
-    port = int(os.environ.get('PORT', 7860))
+    port = int(os.environ.get('PORT', 7860))  # HF Spaces uses port 7860
     app.run(host='0.0.0.0', port=port, debug=False)