Update app.py

app.py

@@ -1,15 +1,15 @@
-from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from pydantic import BaseModel
+from fastapi import FastAPI, HTTPException, BackgroundTasks
 from huggingface_hub import snapshot_download, HfApi
 import os
 import shutil
-from typing import Optional
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
+from dotenv import load_dotenv
 
-app = FastAPI(title="Hugging Face Model Transfer API", version="1.0.0")
-security = HTTPBearer()
+# Load environment variables from .env file
+load_dotenv()
+
+app = FastAPI(title="Hugging Face Model Transfer Service", version="1.0.0")
 
 # Thread pool for running blocking operations
 executor = ThreadPoolExecutor(max_workers=2)
@@ -20,140 +20,100 @@ DOWNLOAD_DIR = "./downloaded_model_data"
 # Ensure the download directory exists
 os.makedirs(DOWNLOAD_DIR, exist_ok=True)
 
-class DownloadRequest(BaseModel):
-    model_repo_id: str
-    download_dir: Optional[str] = DOWNLOAD_DIR
-
-class UploadRequest(BaseModel):
-    dataset_repo_id: str
-    folder_path: str
-    path_in_repo: str
-
-class TransferRequest(BaseModel):
-    model_repo_id: str
-    dataset_repo_id: str
-    path_in_repo: Optional[str] = None
-
-def get_token(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
-    """Extract and validate Hugging Face token from Authorization header."""
-    return credentials.credentials
-
-def download_full_model(repo_id: str, download_dir: str) -> str:
-    """Downloads an entire model from a Hugging Face model repository."""
-    print(f"Downloading full model from {repo_id}...")
-    local_dir = snapshot_download(repo_id=repo_id, cache_dir=download_dir)
-    print(f"Downloaded to: {local_dir}")
-    return local_dir
-
-def upload_folder_to_dataset(dataset_repo_id: str, folder_path: str, path_in_repo: str, token: str):
-    """Uploads a folder to a Hugging Face dataset repository."""
-    api = HfApi(token=token)
-    print(f"Uploading {folder_path} to {dataset_repo_id} at {path_in_repo}...")
-    api.upload_folder(
-        folder_path=folder_path,
-        path_in_repo=path_in_repo,
-        repo_id=dataset_repo_id,
-        repo_type="dataset",
-    )
-    print("Upload complete!")
+# Hardcoded model repository ID
+HARDCODED_MODEL_REPO_ID = "openai/gpt-oss-120b"  # Change this to your desired model
 
-@app.get("/")
-async def root():
-    """Health check endpoint."""
-    return {"message": "Hugging Face Model Transfer API is running"}
+# Hardcoded dataset repository ID
+HARDCODED_DATASET_REPO_ID = "Fred808/helium_memory"  # Change this to your dataset
+
+# Hardcoded path in repository
+HARDCODED_PATH_IN_REPO = "model_data2/"
+
+# Get Hugging Face token from environment variable
+HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is not set")
 
-@app.post("/download")
-async def download_model(request: DownloadRequest, token: str = Depends(get_token)):
-    """Download a model from Hugging Face model repository."""
+def download_full_model() -> str:
+    """Downloads the hardcoded model from Hugging Face model repository."""
+    print(f"Downloading hardcoded model {HARDCODED_MODEL_REPO_ID}...")
     try:
-        # Run the blocking download operation in a thread pool
-        loop = asyncio.get_event_loop()
-        local_dir = await loop.run_in_executor(
-            executor, 
-            download_full_model, 
-            request.model_repo_id, 
-            request.download_dir
+        local_dir = snapshot_download(
+            repo_id=HARDCODED_MODEL_REPO_ID, 
+            cache_dir=DOWNLOAD_DIR,
+            token=HF_TOKEN
         )
-        
-        return {
-            "message": f"Model {request.model_repo_id} downloaded successfully",
-            "local_path": local_dir
-        }
+        print(f"Downloaded to: {local_dir}")
+        return local_dir
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Download failed: {str(e)}")
+        print(f"Download failed: {str(e)}")
+        raise
 
-@app.post("/upload")
-async def upload_folder(request: UploadRequest, token: str = Depends(get_token)):
-    """Upload a folder to a Hugging Face dataset repository."""
+def upload_folder_to_dataset(folder_path: str):
+    """Uploads a folder to the hardcoded Hugging Face dataset repository."""
+    api = HfApi(token=HF_TOKEN)
+    print(f"Uploading {folder_path} to {HARDCODED_DATASET_REPO_ID} at {HARDCODED_PATH_IN_REPO}...")
     try:
-        # Check if folder exists
-        if not os.path.exists(request.folder_path):
-            raise HTTPException(status_code=404, detail=f"Folder not found: {request.folder_path}")
-        
-        # Run the blocking upload operation in a thread pool
-        loop = asyncio.get_event_loop()
-        await loop.run_in_executor(
-            executor,
-            upload_folder_to_dataset,
-            request.dataset_repo_id,
-            request.folder_path,
-            request.path_in_repo,
-            token
+        api.upload_folder(
+            folder_path=folder_path,
+            path_in_repo=HARDCODED_PATH_IN_REPO,
+            repo_id=HARDCODED_DATASET_REPO_ID,
+            repo_type="dataset",
         )
-        
-        return {
-            "message": f"Folder uploaded successfully to {request.dataset_repo_id}",
-            "path_in_repo": request.path_in_repo
-        }
-    except HTTPException:
+        print("Upload complete!")
+    except Exception as e:
+        print(f"Upload failed: {str(e)}")
         raise
+
+def cleanup_download(local_dir: str):
+    """Clean up downloaded files."""
+    try:
+        if os.path.exists(local_dir):
+            shutil.rmtree(local_dir)
+            print(f"Cleaned up: {local_dir}")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Upload failed: {str(e)}")
+        print(f"Cleanup failed: {str(e)}")
 
-@app.post("/transfer")
-async def transfer_model(request: TransferRequest, background_tasks: BackgroundTasks, token: str = Depends(get_token)):
-    """Download a model and upload it to a dataset repository (combined operation)."""
+async def transfer_model():
+    """Download the hardcoded model and upload it to the hardcoded dataset repository."""
     try:
-        # Set default path in repo if not provided
-        path_in_repo = request.path_in_repo or f"model_data/{request.model_repo_id}/"
-        
         # Download the model
         loop = asyncio.get_event_loop()
         local_dir = await loop.run_in_executor(
             executor,
-            download_full_model,
-            request.model_repo_id,
-            DOWNLOAD_DIR
+            download_full_model
         )
         
         # Upload to dataset
         await loop.run_in_executor(
             executor,
             upload_folder_to_dataset,
-            request.dataset_repo_id,
-            local_dir,
-            path_in_repo,
-            token
+            local_dir
         )
         
-        # Clean up downloaded files in background
-        background_tasks.add_task(cleanup_download, local_dir)
+        # Clean up downloaded files
+        cleanup_download(local_dir)
+        
+        print(f"Model {HARDCODED_MODEL_REPO_ID} transferred successfully to {HARDCODED_DATASET_REPO_ID}")
         
-        return {
-            "message": f"Model {request.model_repo_id} transferred successfully to {request.dataset_repo_id}",
-            "path_in_repo": path_in_repo
-        }
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Transfer failed: {str(e)}")
+        print(f"Transfer failed: {str(e)}")
+        raise
 
-def cleanup_download(local_dir: str):
-    """Clean up downloaded files."""
-    try:
-        if os.path.exists(local_dir):
-            shutil.rmtree(local_dir)
-            print(f"Cleaned up: {local_dir}")
-    except Exception as e:
-        print(f"Cleanup failed: {str(e)}")
+@app.get("/")
+async def root():
+    """Health check endpoint."""
+    return {
+        "message": "Hugging Face Model Transfer Service is running", 
+        "hardcoded_model": HARDCODED_MODEL_REPO_ID,
+        "hardcoded_dataset": HARDCODED_DATASET_REPO_ID
+    }
+
+@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():
@@ -167,12 +127,13 @@ async def get_status():
                 "total": disk_usage.total,
                 "used": disk_usage.used,
                 "free": disk_usage.free
-            }
+            },
+            "model": HARDCODED_MODEL_REPO_ID,
+            "dataset": HARDCODED_DATASET_REPO_ID
         }
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Status check failed: {str(e)}")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
- 
+    uvicorn.run(app, host="0.0.0.0", port=7860)