Update app.py

app.py

@@ -9,6 +9,16 @@ import sys
 sys.path.append("/app/ai-toolkit")
 
 from toolkit.job import run_job
+import logging
+import io
+
+log_stream = io.StringIO()
+logger = logging.getLogger("lora_training")
+logger.setLevel(logging.DEBUG)
+handler = logging.StreamHandler(log_stream)
+handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
+logger.addHandler(handler)
+
 
 app = FastAPI()
 
@@ -23,8 +33,10 @@ status = {"running": False, "last_job": None, "error": None}
 def run_lora_training(push_to_hub: bool = False):
     try:
         status.update({"running": True, "error": None})
+        logger.info("Starting training...")
+
         local_dir = f"/tmp/{LORA_NAME}-{uuid.uuid4()}"
-        
+        logger.info(f"Downloading dataset to {local_dir} ...")
         snapshot_download(
             repo_id=REPO_ID,
             repo_type="dataset",
@@ -32,15 +44,27 @@ def run_lora_training(push_to_hub: bool = False):
             local_dir=local_dir,
             local_dir_use_symlinks=False
         )
-        
         training_path = os.path.join(local_dir, FOLDER_IN_REPO)
+        logger.info(f"Building job with training path: {training_path}")
         job = build_job(CONCEPT_SENTENCE, training_path, LORA_NAME, push_to_hub=push_to_hub)
+        logger.info("Running job...")
         run_job(job)
-        
+        logger.info("Training completed successfully.")
         status.update({"running": False, "last_job": job})
     except Exception as e:
+        logger.error(f"Training failed: {e}")
         status.update({"running": False, "error": str(e)})
 
+
+
+
+
+@app.get("/logs")
+def get_logs():
+    return {"logs": log_stream.getvalue()}
+
+
+    
 @app.get("/")
 def root():
     return {"message": "LoRA training FastAPI is live."}