Update app.py

app.py

@@ -1,10 +1,11 @@
-from fastapi import FastAPI, HTTPException, BackgroundTasks
-from huggingface_hub import snapshot_download, HfApi
+from fastapi import FastAPI, HTTPException
+from huggingface_hub import HfApi, hf_hub_download, list_repo_files
 import os
 import shutil
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 from dotenv import load_dotenv
+import tempfile
 
 # Load environment variables from .env file
 load_dotenv()
@@ -14,38 +15,71 @@ app = FastAPI(title="Hugging Face Model Transfer Service", version="1.0.0")
 # Thread pool for running blocking operations
 executor = ThreadPoolExecutor(max_workers=2)
 
-# Local directory to save downloaded model data temporarily
-DOWNLOAD_DIR = "./downloaded_model_data"
-
-# Ensure the download directory exists
-os.makedirs(DOWNLOAD_DIR, exist_ok=True)
+# Get Hugging Face token from environment variable
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is not set")
 
 # Hardcoded model repository ID
-HARDCODED_MODEL_REPO_ID = "openai/gpt-oss-120b"  # Change this to your desired model
+HARDCODED_MODEL_REPO_ID = "bert-base-uncased"  # Change this to your desired model
 
 # Hardcoded dataset repository ID
-HARDCODED_DATASET_REPO_ID = "Fred808/helium_memory"  # Change this to your dataset
+HARDCODED_DATASET_REPO_ID = "your-username/your-dataset-name"  # 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")
-
-def download_full_model() -> str:
-    """Downloads the hardcoded model from Hugging Face model repository."""
-    print(f"Downloading hardcoded model {HARDCODED_MODEL_REPO_ID}...")
+HARDCODED_PATH_IN_REPO = "model_data/"
+
+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}...")
+    
+    # Create a temporary directory for the downloaded files
+    temp_dir = tempfile.mkdtemp(prefix="hf_model_")
+    print(f"Created temporary directory: {temp_dir}")
+    
     try:
-        local_dir = snapshot_download(
-            repo_id=HARDCODED_MODEL_REPO_ID, 
-            cache_dir=DOWNLOAD_DIR,
+        # Get list of all files in the repository
+        api = HfApi(token=HF_TOKEN)
+        files = list_repo_files(
+            repo_id=HARDCODED_MODEL_REPO_ID,
+            repo_type="model",
             token=HF_TOKEN
         )
-        print(f"Downloaded to: {local_dir}")
-        return local_dir
+        
+        print(f"Found {len(files)} files to download")
+        
+        # Download each file
+        for file_path in files:
+            if file_path.endswith('/'):  # Skip directories
+                continue
+                
+            print(f"Downloading: {file_path}")
+            
+            # Create subdirectories if needed
+            file_dir = os.path.join(temp_dir, os.path.dirname(file_path))
+            os.makedirs(file_dir, exist_ok=True)
+            
+            # Download the file
+            local_path = hf_hub_download(
+                repo_id=HARDCODED_MODEL_REPO_ID,
+                filename=file_path,
+                cache_dir=temp_dir,
+                force_download=True,
+                resume_download=False,
+                token=HF_TOKEN
+            )
+            
+            # Move from cache to the desired location
+            final_path = os.path.join(temp_dir, file_path)
+            if local_path != final_path:
+                shutil.move(local_path, final_path)
+        
+        print(f"All files downloaded to: {temp_dir}")
+        return temp_dir
+        
     except Exception as e:
+        # Clean up on error
+        shutil.rmtree(temp_dir, ignore_errors=True)
         print(f"Download failed: {str(e)}")
         raise
 
@@ -59,46 +93,55 @@ def upload_folder_to_dataset(folder_path: str):
             path_in_repo=HARDCODED_PATH_IN_REPO,
             repo_id=HARDCODED_DATASET_REPO_ID,
             repo_type="dataset",
+            token=HF_TOKEN
         )
         print("Upload complete!")
     except Exception as e:
         print(f"Upload failed: {str(e)}")
         raise
 
-def cleanup_download(local_dir: str):
+def cleanup_download(temp_dir: str):
     """Clean up downloaded files."""
     try:
-        if os.path.exists(local_dir):
-            shutil.rmtree(local_dir)
-            print(f"Cleaned up: {local_dir}")
+        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 transfer_model():
     """Download the hardcoded model and upload it to the hardcoded dataset repository."""
+    temp_dir = None
     try:
-        # Download the model
+        # Download the model files
         loop = asyncio.get_event_loop()
-        local_dir = await loop.run_in_executor(
+        temp_dir = await loop.run_in_executor(
             executor,
-            download_full_model
+            download_model_files
         )
         
+        # Verify files were downloaded
+        if not os.listdir(temp_dir):
+            raise Exception("No files were downloaded")
+        
+        print(f"Downloaded files: {os.listdir(temp_dir)}")
+        
         # Upload to dataset
         await loop.run_in_executor(
             executor,
             upload_folder_to_dataset,
-            local_dir
+            temp_dir
         )
         
-        # Clean up downloaded files
-        cleanup_download(local_dir)
-        
         print(f"Model {HARDCODED_MODEL_REPO_ID} transferred successfully to {HARDCODED_DATASET_REPO_ID}")
         
     except Exception as e:
         print(f"Transfer failed: {str(e)}")
         raise
+    finally:
+        # Clean up downloaded files
+        if temp_dir:
+            cleanup_download(temp_dir)
 
 @app.get("/")
 async def root():
@@ -117,19 +160,13 @@ async def startup_event():
 
 @app.get("/status")
 async def get_status():
-    """Get server status and available disk space."""
+    """Get server status."""
     try:
-        disk_usage = shutil.disk_usage(DOWNLOAD_DIR)
         return {
             "status": "healthy",
-            "download_dir": DOWNLOAD_DIR,
-            "disk_space": {
-                "total": disk_usage.total,
-                "used": disk_usage.used,
-                "free": disk_usage.free
-            },
             "model": HARDCODED_MODEL_REPO_ID,
-            "dataset": HARDCODED_DATASET_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)}")