test

Dockerfile

@@ -7,8 +7,21 @@ RUN apt-get update && apt-get install -y \
     git \
     curl \
     build-essential \
+    wget \
     && rm -rf /var/lib/apt/lists/*
 
+# Create cache directory with proper permissions
+RUN mkdir -p /app/cache && chmod 777 /app/cache
+RUN mkdir -p /app/models && chmod 777 /app/models
+
+# Set environment variables for caching
+ENV HF_HOME=/app/cache
+ENV TRANSFORMERS_CACHE=/app/cache
+ENV HF_DATASETS_CACHE=/app/cache
+ENV TORCH_HOME=/app/cache
+ENV PYTHONPATH=/app
+ENV FLASK_APP=app.py
+
 # Copy requirements and install Python dependencies
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
@@ -19,9 +32,5 @@ COPY . .
 # Expose port
 EXPOSE 7860
 
-# Set environment variables
-ENV PYTHONPATH=/app
-ENV FLASK_APP=app.py
-
 # Run the application
 CMD ["python", "app.py"]

app.py

@@ -2,30 +2,11 @@ from flask import Flask, jsonify, request
 import threading
 import time
 import os
-import torch
-from datasets import load_dataset
-from huggingface_hub import login
-from transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig,
-    TrainingArguments,
-    pipeline,
-    logging,
-    DataCollatorForLanguageModeling,
-)
-from peft import (
-    LoraConfig,
-    PeftModel,
-    prepare_model_for_kbit_training,
-    get_peft_model,
-)
-from trl import SFTTrainer, setup_chat_format
+import tempfile
+import shutil
 import uuid
 from datetime import datetime, timedelta
 
-# ============== CONFIGURATION ==============
 app = Flask(__name__)
 
 # Global variables to track training progress
@@ -74,60 +55,78 @@ def train_model_background(job_id):
     progress = training_jobs[job_id]
     
     try:
-        # === Authentication ===
-        import os
+        # 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
         from huggingface_hub import login
-
+        from transformers import (
+            AutoConfig,
+            AutoModelForCausalLM,
+            AutoTokenizer,
+            BitsAndBytesConfig,
+            TrainingArguments,
+            logging,
+        )
+        from peft import (
+            LoraConfig,
+            get_peft_model,
+        )
+        from trl import SFTTrainer, setup_chat_format
+        
+        # === Authentication ===
         hf_token = os.getenv('HF_TOKEN')
-
-        if not hf_token:
-            raise ValueError("HF_TOKEN is not set. Please define it as an environment variable or secret.")
-    
-        login(token=hf_token)
-
+        if hf_token:
+            login(token=hf_token)
         
         progress.status = "loading_model"
         progress.message = "Loading base model and tokenizer..."
 
         # === Configuration ===
-        base_model = "meta-llama/Llama-3.2-1B"
+        base_model = "microsoft/DialoGPT-small"  # Smaller model for testing
         dataset_name = "ruslanmv/ai-medical-chatbot"
-        new_model = f"Llama-3.2-3B-chat-doctor-{job_id}"
-        torch_dtype = torch.float16
-        attn_implementation = "eager"
-
-        # === QLoRA Config ===
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch_dtype,
-            bnb_4bit_use_double_quant=True,
-        )
-
-        # === Load Model and Tokenizer ===
+        new_model = f"trained-model-{job_id}"
+        
+        # === Load Model and Tokenizer (without quantization for simplicity) ===
         model = AutoModelForCausalLM.from_pretrained(
             base_model,
-            quantization_config=bnb_config,
-            device_map="auto",
-            attn_implementation=attn_implementation
+            cache_dir=temp_dir,
+            torch_dtype=torch.float32,  # Use float32 for compatibility
+            device_map="auto" if torch.cuda.is_available() else "cpu",
+            trust_remote_code=True
         )
-        tokenizer = AutoTokenizer.from_pretrained(base_model)
-        model, tokenizer = setup_chat_format(model, tokenizer)
+        
+        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
 
         progress.status = "preparing_model"
         progress.message = "Setting up LoRA configuration..."
 
-        # === LoRA Config ===
+        # === LoRA Config (simplified) ===
         peft_config = LoraConfig(
-            r=16,
-            lora_alpha=32,
-            lora_dropout=0.05,
+            r=8,  # Smaller rank
+            lora_alpha=16,
+            lora_dropout=0.1,
             bias="none",
             task_type="CAUSAL_LM",
-            target_modules=[
-                'up_proj', 'down_proj', 'gate_proj',
-                'k_proj', 'q_proj', 'v_proj', 'o_proj'
-            ]
         )
         model = get_peft_model(model, peft_config)
 
@@ -135,29 +134,26 @@ def train_model_background(job_id):
         progress.message = "Loading and preparing dataset..."
 
         # === Load & Prepare Dataset ===
-        dataset = load_dataset(dataset_name, split="all")
-        dataset = dataset.shuffle(seed=65).select(range(1000))  # Use 1000 samples
-
-        def format_chat_template(row, tokenizer):
-            row_json = [
-                {"role": "user", "content": row["Patient"]},
-                {"role": "assistant", "content": row["Doctor"]}
-            ]
-            row["text"] = tokenizer.apply_chat_template(row_json, tokenize=False)
-            return row
-
-        dataset = dataset.map(
-            format_chat_template,
-            fn_kwargs={"tokenizer": tokenizer},
-            num_proc=4
+        dataset = load_dataset(
+            dataset_name, 
+            split="all",
+            cache_dir=temp_dir,
+            trust_remote_code=True
         )
+        dataset = dataset.shuffle(seed=65).select(range(100))  # Use only 100 samples for testing
+
+        def format_chat_template(row):
+            # Simple formatting without chat template
+            text = f"Patient: {row['Patient']}\nDoctor: {row['Doctor']}"
+            return {"text": text}
 
+        dataset = dataset.map(format_chat_template, num_proc=1)
         dataset = dataset.train_test_split(test_size=0.1)
 
         # Calculate total training steps
         train_size = len(dataset["train"])
         batch_size = 1
-        gradient_accumulation_steps = 2
+        gradient_accumulation_steps = 1
         num_epochs = 1
         
         steps_per_epoch = train_size // (batch_size * gradient_accumulation_steps)
@@ -168,29 +164,33 @@ def train_model_background(job_id):
         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=new_model,
+            output_dir=output_dir,
             per_device_train_batch_size=batch_size,
             per_device_eval_batch_size=1,
             gradient_accumulation_steps=gradient_accumulation_steps,
-            optim="paged_adamw_32bit",
+            optim="adamw_torch",  # Use standard optimizer
             num_train_epochs=num_epochs,
-            eval_steps=0.2,
+            eval_steps=0.5,
             logging_steps=1,
-            warmup_steps=10,
+            warmup_steps=5,
             logging_strategy="steps",
-            learning_rate=2e-5,
+            learning_rate=5e-5,
             fp16=False,
             bf16=False,
             group_by_length=True,
-            save_steps=50,
-            save_total_limit=2,
-            report_to=None  # Disable wandb for HF Spaces
+            save_steps=10,
+            save_total_limit=1,
+            report_to=None,
+            dataloader_num_workers=0,
+            remove_unused_columns=False,
+            load_best_model_at_end=False,
+            evaluation_strategy="no"  # Disable evaluation for simplicity
         )
 
-        # === Data Collator ===
-        tokenizer.model_max_length = 512
-        
         # Custom callback to track progress
         class ProgressCallback:
             def __init__(self, progress_tracker):
@@ -200,7 +200,7 @@ def train_model_background(job_id):
             def on_log(self, args, state, control, model=None, logs=None, **kwargs):
                 current_time = time.time()
                 # Update every 10 seconds or on significant step changes
-                if current_time - self.last_update >= 10 or state.global_step % 10 == 0:
+                if current_time - self.last_update >= 10 or state.global_step % 5 == 0:
                     self.progress_tracker.update_progress(
                         state.global_step, 
                         state.max_steps,
@@ -212,24 +212,44 @@ def train_model_background(job_id):
         trainer = SFTTrainer(
             model=model,
             train_dataset=dataset["train"],
-            eval_dataset=dataset["test"],
             peft_config=peft_config,
             args=training_args,
-            callbacks=[ProgressCallback(progress)]
+            callbacks=[ProgressCallback(progress)],
+            tokenizer=tokenizer,
+            max_seq_length=256,  # Shorter sequences
         )
 
         # === Train & Save ===
         trainer.train()
-        trainer.save_model(new_model)
+        trainer.save_model(output_dir)
         
         progress.status = "completed"
         progress.progress = 100
-        progress.message = f"Training completed! Model saved as {new_model}"
+        progress.message = f"Training completed! Model saved to {output_dir}"
+        
+        # Clean up temporary directory after a delay
+        def cleanup_temp_dir():
+            time.sleep(300)  # Wait 5 minutes before cleanup
+            try:
+                shutil.rmtree(temp_dir)
+            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
 
 # ============== API ROUTES ==============
 @app.route('/api/train', methods=['POST'])