Update app.py

app.py

@@ -102,15 +102,42 @@ def search_hub(query: str, search_type: str, token: str = None) -> pd.DataFrame:
 async def download_readme(session: aiohttp.ClientSession, item: Dict, token: str) -> tuple[str, str]:
     """Download README.md file for a given item."""
     item_id = item['id']
-    raw_url = f"https://huggingface.co/{item_id}/raw/main/README.md"
+    
+    # Different base URLs for different repository types
+    if 'datasets' in item['link']:
+        raw_url = f"https://huggingface.co/datasets/{item_id}/raw/main/README.md"
+        alt_url = f"https://huggingface.co/datasets/{item_id}/raw/master/README.md"
+    elif 'spaces' in item['link']:
+        raw_url = f"https://huggingface.co/spaces/{item_id}/raw/main/README.md"
+        alt_url = f"https://huggingface.co/spaces/{item_id}/raw/master/README.md"
+    else:  # Models
+        raw_url = f"https://huggingface.co/{item_id}/raw/main/README.md"
+        alt_url = f"https://huggingface.co/{item_id}/raw/master/README.md"
+    
     headers = {"Authorization": f"Bearer {token}"} if token else {}
     
     try:
+        # Try main branch first
         async with session.get(raw_url, headers=headers) as response:
             if response.status == 200:
                 content = await response.text()
                 return item_id.replace('/', '_'), content
-            return item_id.replace('/', '_'), f"# Error downloading README for {item_id}\nStatus code: {response.status}"
+                
+            # If main branch fails, try master branch
+            if response.status in [401, 404]:
+                async with session.get(alt_url, headers=headers) as alt_response:
+                    if alt_response.status == 200:
+                        content = await alt_response.text()
+                        return item_id.replace('/', '_'), content
+                    
+            # If both attempts fail, return error message
+            error_msg = f"# Error downloading README for {item_id}\n"
+            if response.status == 401:
+                error_msg += "Authentication required. Please provide a valid HuggingFace token."
+            else:
+                error_msg += f"Status code: {response.status}"
+            return item_id.replace('/', '_'), error_msg
+            
     except Exception as e:
         return item_id.replace('/', '_'), f"# Error downloading README for {item_id}\nError: {str(e)}"
 
@@ -118,9 +145,10 @@ async def download_all_readmes(data: List[Dict], token: str) -> tuple[str, str]:
     """Download all README files and create a zip archive."""
     if not data:
         return "", "No results to download"
-        
+    
     zip_buffer = io.BytesIO()
     status_message = "Downloading READMEs..."
+    failed_downloads = []
     
     async with aiohttp.ClientSession() as session:
         tasks = [download_readme(session, item, token) for item in data]
@@ -128,11 +156,17 @@ async def download_all_readmes(data: List[Dict], token: str) -> tuple[str, str]:
         
         with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
             for filename, content in results:
+                if "Error downloading README" in content:
+                    failed_downloads.append(filename)
                 zip_file.writestr(f"{filename}.md", content)
     
     zip_buffer.seek(0)
     base64_zip = base64.b64encode(zip_buffer.getvalue()).decode()
     
+    status = "READMEs ready for download!"
+    if failed_downloads:
+        status += f" (Failed to download {len(failed_downloads)} READMEs)"
+    
     download_link = f"""
     <div style="margin-top: 10px;">
         <a href="data:application/zip;base64,{base64_zip}" 
@@ -142,10 +176,11 @@ async def download_all_readmes(data: List[Dict], token: str) -> tuple[str, str]:
                   text-decoration: none; border-radius: 5px;">
             📥 Download READMEs Archive
         </a>
+        {f'<p style="color: #ff6b6b; margin-top: 10px;">Note: Some READMEs could not be downloaded. Please check the zip file for details.</p>' if failed_downloads else ''}
     </div>
     """
     
-    return download_link, "READMEs ready for download!"
+    return download_link, status
 
 def download_repository(repo_id: str, repo_type: str, temp_dir: str, token: str) -> str:
     """Download a single repository."""