Update app.py

app.py

@@ -5,9 +5,18 @@ import subprocess
 import os
 import uuid
 from huggingface_hub import HfApi, HfFolder
+import logging
 
 app = FastAPI()
 
+# Configure logging
+logging.basicConfig(
+    filename='training.log',
+    filemode='a',
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    level=logging.INFO
+)
+
 # Define the expected payload structure
 class TrainingRequest(BaseModel):
     task: str  # 'generation' or 'classification'
@@ -24,9 +33,17 @@ if not HF_API_TOKEN:
 HfFolder.save_token(HF_API_TOKEN)
 api = HfApi()
 
+@app.get("/")
+def read_root():
+    return {
+        "message": "Welcome to the Training Space API!",
+        "instructions": "To train a model, send a POST request to /train with the required parameters."
+    }
+
 @app.post("/train")
 def train_model(request: TrainingRequest):
     try:
+        logging.info(f"Received training request for model: {request.model_name}, Task: {request.task}")
         # Create a unique directory for this training session
         session_id = str(uuid.uuid4())
         session_dir = f"./training_sessions/{session_id}"
@@ -53,7 +70,23 @@ def train_model(request: TrainingRequest):
         # Start the training process as a background task
         subprocess.Popen(cmd, cwd=session_dir)
         
+        logging.info(f"Training started for model: {request.model_name}, Session ID: {session_id}")
+        
         return {"status": "Training started", "session_id": session_id}
     
     except Exception as e:
+        logging.error(f"Error during training request: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
+
+# Optional: Status Endpoint
+@app.get("/status/{session_id}")
+def get_status(session_id: str):
+    session_dir = f"./training_sessions/{session_id}"
+    log_file = os.path.join(session_dir, "training.log")
+    if not os.path.exists(log_file):
+        raise HTTPException(status_code=404, detail="Session ID not found.")
+    
+    with open(log_file, "r", encoding="utf-8") as f:
+        logs = f.read()
+    
+    return {"session_id": session_id, "logs": logs}