Update app.py

app.py

@@ -29,6 +29,10 @@ HARDCODED_DATASET_REPO_ID = "Fred808/helium_memory"  # Change this to your datas
 # Hardcoded path in repository
 HARDCODED_PATH_IN_REPO = "model_data/"
 
+# Track transfer status
+transfer_completed = False
+transfer_error = None
+
 def download_model_files() -> str:
     """Download all files from a model repository to a temporary directory."""
     print(f"Downloading all files from model {HARDCODED_MODEL_REPO_ID}...")
@@ -109,8 +113,9 @@ def cleanup_download(temp_dir: str):
     except Exception as e:
         print(f"Cleanup failed: {str(e)}")
 
-async def transfer_model():
-    """Download the hardcoded model and upload it to the hardcoded dataset repository."""
+async def run_transfer():
+    """Run the transfer process and update status."""
+    global transfer_completed, transfer_error
     temp_dir = None
     try:
         # Download the model files
@@ -134,9 +139,12 @@ async def transfer_model():
         )
         
         print(f"Model {HARDCODED_MODEL_REPO_ID} transferred successfully to {HARDCODED_DATASET_REPO_ID}")
+        transfer_completed = True
         
     except Exception as e:
-        print(f"Transfer failed: {str(e)}")
+        error_msg = f"Transfer failed: {str(e)}"
+        print(error_msg)
+        transfer_error = error_msg
         raise
     finally:
         # Clean up downloaded files
@@ -146,31 +154,35 @@ async def transfer_model():
 @app.get("/")
 async def root():
     """Health check endpoint."""
+    status = "completed" if transfer_completed else "running" if transfer_error is None else "failed"
     return {
         "message": "Hugging Face Model Transfer Service is running", 
-        "hardcoded_model": HARDCODED_MODEL_REPO_ID,
-        "hardcoded_dataset": HARDCODED_DATASET_REPO_ID
+        "status": status,
+        "model": HARDCODED_MODEL_REPO_ID,
+        "dataset": HARDCODED_DATASET_REPO_ID,
+        "error": transfer_error
+    }
+
+@app.get("/status")
+async def get_status():
+    """Get transfer status."""
+    status = "completed" if transfer_completed else "running" if transfer_error is None else "failed"
+    return {
+        "status": status,
+        "model": HARDCODED_MODEL_REPO_ID,
+        "dataset": HARDCODED_DATASET_REPO_ID,
+        "path_in_repo": HARDCODED_PATH_IN_REPO,
+        "error": transfer_error
     }
 
 @app.on_event("startup")
 async def startup_event():
     """Run the transfer process when the application starts."""
     print("Starting model transfer process...")
-    asyncio.create_task(transfer_model())
-
-@app.get("/status")
-async def get_status():
-    """Get server status."""
-    try:
-        return {
-            "status": "healthy",
-            "model": HARDCODED_MODEL_REPO_ID,
-            "dataset": HARDCODED_DATASET_REPO_ID,
-            "path_in_repo": HARDCODED_PATH_IN_REPO
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Status check failed: {str(e)}")
+    # Run transfer in background without waiting for completion
+    asyncio.create_task(run_transfer())
 
 if __name__ == "__main__":
     import uvicorn
+    # Run the server indefinitely
     uvicorn.run(app, host="0.0.0.0", port=7860)