Update app.py

app.py

@@ -184,11 +184,11 @@ async def root():
 async def health_check():
     return {"status": "healthy"}
 
-@app.get("/v1/bert/training-status")
+@app.get("/training-status")
 async def get_training_status():
     return training_status
 
-@app.post("/v1/bert/upload")
+@app.post("/upload")
 async def upload_file(file: UploadFile = File(...)):
     """Upload a CSV file for training or validation"""
     if not file.filename.endswith('.csv'):
@@ -200,7 +200,7 @@ async def upload_file(file: UploadFile = File(...)):
     
     return {"message": f"File {file.filename} uploaded successfully", "file_path": str(file_path)}
 
-@app.post("/v1/bert/train", response_model=TrainingResponse)
+@app.post("/bert/train", response_model=TrainingResponse)
 async def start_training(
     config: TrainingConfig,
     background_tasks: BackgroundTasks,
@@ -224,7 +224,7 @@ async def start_training(
     
     background_tasks.add_task(train_model_task, config, file_path, training_id)
     
-    download_url = f"/v1/bert/download-model/{training_id}"
+    download_url = f"/bert/download-model/{training_id}"
     
     return TrainingResponse(
         message="Training started successfully",
@@ -233,7 +233,7 @@ async def start_training(
         download_url=download_url
     )
 
-@app.post("/v1/bert/validate")
+@app.post("/bert/validate")
 async def validate_model(
     file: UploadFile = File(...),
     model_name: str = "BERT_model"
@@ -319,7 +319,7 @@ async def validate_model(
         if os.path.exists(file_path):
             os.remove(file_path)
 
-@app.post("/v1/bert/predict")
+@app.post("/bert/predict")
 async def predict(
     request: Optional[PredictionRequest] = None,
     file: Optional[UploadFile] = File(None),
@@ -510,80 +510,51 @@ async def train_model_task(config: TrainingConfig, file_path: str, training_id:
         data_df_original, label_encoders = load_and_preprocess_data(file_path)
         save_label_encoders(label_encoders)
         
-        train_df, val_df = train_test_split(
-            data_df_original,
-            test_size=config.test_size,
-            random_state=config.random_state,
-            stratify=data_df_original[LABEL_COLUMNS[0]]
-        )
-        
-        train_texts = train_df[TEXT_COLUMN]
-        val_texts = val_df[TEXT_COLUMN]
-        train_labels_array = train_df[LABEL_COLUMNS].values
-        val_labels_array = val_df[LABEL_COLUMNS].values
+        texts = data_df_original[TEXT_COLUMN]
+        labels_array = data_df_original[LABEL_COLUMNS].values
         
-        train_metadata_df = train_df[METADATA_COLUMNS] if METADATA_COLUMNS and all(col in train_df.columns for col in METADATA_COLUMNS) else None
-        val_metadata_df = val_df[METADATA_COLUMNS] if METADATA_COLUMNS and all(col in val_df.columns for col in METADATA_COLUMNS) else None
+        metadata_df = data_df_original[METADATA_COLUMNS] if METADATA_COLUMNS and all(col in data_df_original.columns for col in METADATA_COLUMNS) else None
         
         num_labels_list = get_num_labels(label_encoders)
         tokenizer = get_tokenizer(config.model_name)
         
-        if train_metadata_df is not None and val_metadata_df is not None:
-            metadata_dim = train_metadata_df.shape[1]
-            train_dataset = ComplianceDatasetWithMetadata(
-                train_texts.tolist(),
-                train_metadata_df.values,
-                train_labels_array,
-                tokenizer,
-                config.max_length
-            )
-            val_dataset = ComplianceDatasetWithMetadata(
-                val_texts.tolist(),
-                val_metadata_df.values,
-                val_labels_array,
+        if metadata_df is not None:
+            metadata_dim = metadata_df.shape[1]
+            dataset = ComplianceDatasetWithMetadata(
+                texts.tolist(),
+                metadata_df.values,
+                labels_array,
                 tokenizer,
                 config.max_length
             )
             model = BertMultiOutputModel(num_labels_list, metadata_dim=metadata_dim).to(DEVICE)
         else:
-            train_dataset = ComplianceDataset(
-                train_texts.tolist(),
-                train_labels_array,
-                tokenizer,
-                config.max_length
-            )
-            val_dataset = ComplianceDataset(
-                val_texts.tolist(),
-                val_labels_array,
+            dataset = ComplianceDataset(
+                texts.tolist(),
+                labels_array,
                 tokenizer,
                 config.max_length
             )
             model = BertMultiOutputModel(num_labels_list).to(DEVICE)
         
-        train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
-        val_loader = DataLoader(val_dataset, batch_size=config.batch_size)
+        train_loader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)
         
         criterions = initialize_criterions(num_labels_list)
         optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
         
-        best_val_loss = float('inf')
         for epoch in range(config.num_epochs):
             training_status["current_epoch"] = epoch + 1
             
             train_loss = train_model(model, train_loader, criterions, optimizer)
-            val_metrics, _, _ = evaluate_model(model, val_loader)
-            
             training_status["current_loss"] = train_loss
             
-            if val_metrics["loss"] < best_val_loss:
-                best_val_loss = val_metrics["loss"]
-                save_model(model, training_id)
+            # Save model after each epoch
+            save_model(model, training_id)
         
         training_status.update({
             "is_training": False,
             "end_time": datetime.now().isoformat(),
-            "status": "completed",
-            "metrics": summarize_metrics(val_metrics).to_dict()
+            "status": "completed"
         })
         
     except Exception as e: