Update app.py

app.py

@@ -17,9 +17,9 @@ app = FastAPI(title="Hugging Face Model Transfer Service", version="1.0.0")
 executor = ThreadPoolExecutor(max_workers=2)
 
 # Get Hugging Face token from environment variable
-HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 if not HF_TOKEN:
-    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
+    raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is not set")
 
 # Hardcoded model repository ID
 HARDCODED_MODEL_REPO_ID = "openai/gpt-oss-20b"  # Your model
@@ -28,12 +28,13 @@ HARDCODED_MODEL_REPO_ID = "openai/gpt-oss-20b"  # Your model
 HARDCODED_DATASET_REPO_ID = "Fred808/helium_memory"  # Your dataset
 
 # Hardcoded path in repository
-HARDCODED_PATH_IN_REPO = "model_data/gpt-oss-20b"
+HARDCODED_PATH_IN_REPO = "model_data/"
 
 # Track transfer status
 transfer_completed = False
 transfer_error = None
 transfer_started = False
+temp_dir = None
 
 def download_model_files() -> str:
     """Download all files from a model repository to a temporary directory."""
@@ -125,20 +126,10 @@ def upload_folder_to_dataset(folder_path: str):
         print(f"Upload failed: {str(e)}")
         raise
 
-def cleanup_download(temp_dir: str):
-    """Clean up downloaded files."""
-    try:
-        if os.path.exists(temp_dir):
-            shutil.rmtree(temp_dir)
-            print(f"Cleaned up temporary directory: {temp_dir}")
-    except Exception as e:
-        print(f"Cleanup failed: {str(e)}")
-
 async def run_transfer():
     """Run the transfer process and update status."""
-    global transfer_completed, transfer_error, transfer_started
+    global transfer_completed, transfer_error, transfer_started, temp_dir
     transfer_started = True
-    temp_dir = None
     try:
         # Download the model files
         loop = asyncio.get_event_loop()
@@ -167,10 +158,6 @@ async def run_transfer():
         error_msg = f"Transfer failed: {str(e)}"
         print(error_msg)
         transfer_error = error_msg
-    finally:
-        # Clean up downloaded files
-        if temp_dir:
-            cleanup_download(temp_dir)
 
 @app.get("/")
 async def root():
@@ -181,7 +168,8 @@ async def root():
         "status": status,
         "model": HARDCODED_MODEL_REPO_ID,
         "dataset": HARDCODED_DATASET_REPO_ID,
-        "error": transfer_error
+        "error": transfer_error,
+        "temp_dir": temp_dir if temp_dir else "Not created yet"
     }
 
 @app.get("/status")
@@ -193,9 +181,25 @@ async def get_status():
         "model": HARDCODED_MODEL_REPO_ID,
         "dataset": HARDCODED_DATASET_REPO_ID,
         "path_in_repo": HARDCODED_PATH_IN_REPO,
-        "error": transfer_error
+        "error": transfer_error,
+        "temp_dir": temp_dir if temp_dir else "Not created yet"
     }
 
+@app.get("/cleanup")
+async def cleanup():
+    """Manual cleanup endpoint to remove downloaded files."""
+    global temp_dir
+    if temp_dir and os.path.exists(temp_dir):
+        try:
+            shutil.rmtree(temp_dir)
+            message = f"Cleaned up temporary directory: {temp_dir}"
+            temp_dir = None
+            return {"message": message}
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Cleanup failed: {str(e)}")
+    else:
+        return {"message": "No temporary directory to clean up"}
+
 @app.on_event("startup")
 async def startup_event():
     """Run the transfer process when the application starts."""