fixing training

app.py

@@ -6,6 +6,9 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 from typing import List, Optional
+from datasets import load_dataset
+from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling
+import json
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -31,6 +34,26 @@ class HealthResponse(BaseModel):
     gpu_available: bool
     device: str
 
+class TrainRequest(BaseModel):
+    dataset_path: str
+    num_epochs: Optional[int] = 3
+    batch_size: Optional[int] = 4
+    learning_rate: Optional[float] = 2e-5
+
+class TrainResponse(BaseModel):
+    status: str
+    message: str
+
+# Add training status tracking
+class TrainingStatus:
+    def __init__(self):
+        self.is_training = False
+        self.current_epoch = 0
+        self.current_loss = None
+        self.status = "idle"
+
+training_status = TrainingStatus()
+
 # Initialize FastAPI app
 app = FastAPI(
     title="Medical LLaMA API",
@@ -133,4 +156,166 @@ async def startup_event():
         tokenizer, model = init_model()
         logger.info("Model loaded successfully")
     except Exception as e:
-        logger.error(f"Failed to load model: {str(e)}") 
+        logger.error(f"Failed to load model: {str(e)}")
+
+@app.post("/train", response_model=TrainResponse, tags=["Training"])
+async def train_model(request: TrainRequest, background_tasks: BackgroundTasks):
+    """
+    Start model training with the specified dataset
+    
+    Parameters:
+    - dataset_path: Path to the JSON dataset file
+    - num_epochs: Number of training epochs
+    - batch_size: Training batch size
+    - learning_rate: Learning rate for training
+    """
+    if training_status.is_training:
+        raise HTTPException(status_code=400, detail="Training is already in progress")
+    
+    try:
+        # Verify dataset exists
+        if not os.path.exists(request.dataset_path):
+            raise HTTPException(status_code=404, detail="Dataset file not found")
+        
+        # Start training in background
+        background_tasks.add_task(
+            run_training,
+            request.dataset_path,
+            request.num_epochs,
+            request.batch_size,
+            request.learning_rate
+        )
+        
+        return TrainResponse(
+            status="started",
+            message="Training started in background"
+        )
+        
+    except Exception as e:
+        logger.error(f"Training setup error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/train/status", tags=["Training"])
+async def get_training_status():
+    """
+    Get current training status
+    """
+    return {
+        "is_training": training_status.is_training,
+        "current_epoch": training_status.current_epoch,
+        "current_loss": training_status.current_loss,
+        "status": training_status.status
+    }
+
+# Add training function
+async def run_training(dataset_path: str, num_epochs: int, batch_size: int, learning_rate: float):
+    global model, tokenizer, training_status
+    
+    try:
+        training_status.is_training = True
+        training_status.status = "loading_dataset"
+        
+        # Load dataset
+        dataset = load_dataset("json", data_files=dataset_path)
+        
+        training_status.status = "preprocessing"
+        
+        # Preprocess function
+        def preprocess_function(examples):
+            return tokenizer(
+                examples["text"],
+                truncation=True,
+                padding="max_length",
+                max_length=512
+            )
+        
+        # Tokenize dataset
+        tokenized_dataset = dataset.map(
+            preprocess_function,
+            batched=True,
+            remove_columns=dataset["train"].column_names
+        )
+        
+        training_status.status = "training"
+        
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir=f"{model_output_path}/checkpoints",
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=4,
+            num_train_epochs=num_epochs,
+            learning_rate=learning_rate,
+            fp16=True,
+            save_steps=500,
+            logging_steps=100,
+        )
+        
+        # Initialize trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset["train"],
+            data_collator=DataCollatorForLanguageModeling(
+                tokenizer=tokenizer,
+                mlm=False
+            ),
+        )
+        
+        # Training callback to update status
+        class TrainingCallback(trainer.callback_handler):
+            def on_epoch_begin(self, args, state, control, **kwargs):
+                training_status.current_epoch = state.epoch
+            
+            def on_log(self, args, state, control, logs=None, **kwargs):
+                if logs:
+                    training_status.current_loss = logs.get("loss", None)
+        
+        trainer.add_callback(TrainingCallback)
+        
+        # Start training
+        trainer.train()
+        
+        # Save the model
+        training_status.status = "saving"
+        model.save_pretrained(model_output_path)
+        tokenizer.save_pretrained(model_output_path)
+        
+        training_status.status = "completed"
+        logger.info("Training completed successfully")
+        
+    except Exception as e:
+        training_status.status = f"failed: {str(e)}"
+        logger.error(f"Training error: {str(e)}")
+        raise
+    
+    finally:
+        training_status.is_training = False
+
+# Update model initialization
+def init_model():
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Loading model on device: {device}")
+        
+        # Try to load fine-tuned model if it exists
+        if os.path.exists(model_output_path):
+            tokenizer = AutoTokenizer.from_pretrained(model_output_path)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_output_path,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                device_map="auto"
+            )
+        else:
+            # Load base model if no fine-tuned model exists
+            model_name = "nvidia/Meta-Llama-3.2-3B-Instruct-ONNX-INT4"
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                device_map="auto"
+            )
+        
+        return tokenizer, model
+    except Exception as e:
+        logger.error(f"Model initialization error: {str(e)}")
+        raise