Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File
+from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, Form
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
 from typing import Optional, Dict, Any, List
@@ -18,6 +18,7 @@ import zipfile
 import io
 import numpy as np
 import sys
+import json
 
 
 # Import existing utilities
@@ -88,7 +89,6 @@ class TrainingConfig(BaseModel):
     learning_rate: float = 2e-5
     num_epochs: int = 2
     max_length: int = 128
-    test_size: float = 0.2
     random_state: int = 42
 
 class TrainingResponse(BaseModel):
@@ -190,8 +190,8 @@ async def get_training_status():
 
 @app.post("/v1/bert/train", response_model=TrainingResponse)
 async def start_training(
-    config: TrainingConfig,
-    background_tasks: BackgroundTasks,
+    config: str = Form(...),
+    background_tasks: BackgroundTasks = None,
     file: UploadFile = File(...)
 ):
     if training_status["is_training"]:
@@ -200,6 +200,15 @@ async def start_training(
     if not file.filename.endswith('.csv'):
         raise HTTPException(status_code=400, detail="Only CSV files are allowed")
     
+    try:
+        # Parse the config JSON string into a TrainingConfig object
+        config_dict = json.loads(config)
+        training_config = TrainingConfig(**config_dict)
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid config JSON format")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid config parameters: {str(e)}")
+    
     file_path = UPLOAD_DIR / file.filename
     with file_path.open("wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
@@ -209,12 +218,12 @@ async def start_training(
     training_status.update({
         "is_training": True,
         "current_epoch": 0,
-        "total_epochs": config.num_epochs,
+        "total_epochs": training_config.num_epochs,
         "start_time": datetime.now().isoformat(),
         "status": "starting"
     })
     
-    background_tasks.add_task(train_model_task, config, str(file_path), training_id)
+    background_tasks.add_task(train_model_task, training_config, str(file_path), training_id)
     
     download_url = f"/v1/bert/download-model/{training_id}"
     
@@ -314,7 +323,7 @@ async def validate_model(
 @app.post("/v1/bert/predict")
 async def predict(
     request: Optional[PredictionRequest] = None,
-    file: Optional[UploadFile] = File(None),
+    file: UploadFile = File(None),
     model_name: str = "BERT_model"
 ):
     """
@@ -325,6 +334,7 @@ async def predict(
     2. Upload a CSV file with multiple transactions
     
     Parameters:
+    - file: CSV file containing transactions for batch prediction
     - model_name: Name of the model to use for prediction (default: "BERT_model")
     """
     try:
@@ -339,7 +349,7 @@ async def predict(
         model.eval()
         
         # Handle batch prediction from CSV
-        if file is not None and file.filename:
+        if file and file.filename:
             if not file.filename.endswith('.csv'):
                 raise HTTPException(status_code=400, detail="Only CSV files are allowed")
             
@@ -398,7 +408,7 @@ async def predict(
                     os.remove(file_path)
         
         # Handle single prediction
-        elif request is not None and request.transaction_data:
+        elif request and request.transaction_data:
             input_data = pd.DataFrame([request.transaction_data.dict()])
             
             text_input = f"""
@@ -430,20 +440,6 @@ async def predict(
             Sanction Description: {input_data['Sanction_Description'].iloc[0]}
             Checker Notes: {input_data['Checker_Notes'].iloc[0]}
             Sanction Context: {input_data['Sanction_Context'].iloc[0]}
-            Maker Action: {input_data['Maker_Action'].iloc[0]}
-            Customer Type: {input_data['Customer_Type'].iloc[0]}
-            Industry: {input_data['Industry'].iloc[0]}
-            Transaction Type: {input_data['Transaction_Type'].iloc[0]}
-            Transaction Channel: {input_data['Transaction_Channel'].iloc[0]}
-            Geographic Origin: {input_data['Geographic_Origin'].iloc[0]}
-            Geographic Destination: {input_data['Geographic_Destination'].iloc[0]}
-            Risk Category: {input_data['Risk_Category'].iloc[0]}
-            Risk Drivers: {input_data['Risk_Drivers'].iloc[0]}
-            Alert Status: {input_data['Alert_Status'].iloc[0]}
-            Investigation Outcome: {input_data['Investigation_Outcome'].iloc[0]}
-            Source of Funds: {input_data['Source_Of_Funds'].iloc[0]}
-            Purpose of Transaction: {input_data['Purpose_Of_Transaction'].iloc[0]}
-            Beneficial Owner: {input_data['Beneficial_Owner'].iloc[0]}
             """
             
             dataset = ComplianceDataset(