Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,273 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import List, Dict
|
2 |
+
import httpx
|
3 |
+
import gradio as gr
|
4 |
+
import pandas as pd
|
5 |
+
from huggingface_hub import HfApi, ModelCard, snapshot_download
|
6 |
+
import base64
|
7 |
+
import io
|
8 |
+
import zipfile
|
9 |
+
import asyncio
|
10 |
+
import aiohttp
|
11 |
+
from pathlib import Path
|
12 |
+
import emoji
|
13 |
+
import tempfile
|
14 |
+
import shutil
|
15 |
+
import os
|
16 |
+
|
17 |
+
def search_hub(query: str, search_type: str) -> pd.DataFrame:
|
18 |
+
api = HfApi()
|
19 |
+
if search_type == "Models":
|
20 |
+
results = api.list_models(search=query)
|
21 |
+
data = [{"id": model.modelId, "author": model.author, "downloads": model.downloads, "link": f"https://huggingface.co/{model.modelId}"} for model in results]
|
22 |
+
elif search_type == "Datasets":
|
23 |
+
results = api.list_datasets(search=query)
|
24 |
+
data = [{"id": dataset.id, "author": dataset.author, "downloads": dataset.downloads, "link": f"https://huggingface.co/datasets/{dataset.id}"} for dataset in results]
|
25 |
+
elif search_type == "Spaces":
|
26 |
+
results = api.list_spaces(search=query)
|
27 |
+
data = [{"id": space.id, "author": space.author, "link": f"https://huggingface.co/spaces/{space.id}"} for space in results]
|
28 |
+
else:
|
29 |
+
data = []
|
30 |
+
|
31 |
+
for i, item in enumerate(data, 1):
|
32 |
+
item['number'] = i
|
33 |
+
item['formatted_link'] = format_link(item, i, search_type)
|
34 |
+
|
35 |
+
return pd.DataFrame(data)
|
36 |
+
|
37 |
+
def format_link(item: Dict, number: int, search_type: str) -> str:
|
38 |
+
link = item['link']
|
39 |
+
readme_link = f"{link}/blob/main/README.md"
|
40 |
+
title = f"{number}. {item['id']}"
|
41 |
+
|
42 |
+
metadata = f"Author: {item['author']}"
|
43 |
+
if 'downloads' in item:
|
44 |
+
metadata += f", Downloads: {item['downloads']}"
|
45 |
+
|
46 |
+
html = f"""
|
47 |
+
<div style="margin-bottom: 10px;">
|
48 |
+
<strong>{title}</strong><br>
|
49 |
+
<a href="{link}" target="_blank" style="color: #4a90e2; text-decoration: none;">View {search_type[:-1]}</a> |
|
50 |
+
<a href="{readme_link}" target="_blank" style="color: #4a90e2; text-decoration: none;">View README</a><br>
|
51 |
+
<small>{metadata}</small>
|
52 |
+
</div>
|
53 |
+
"""
|
54 |
+
return html
|
55 |
+
|
56 |
+
async def download_readme(session: aiohttp.ClientSession, item: Dict) -> tuple[str, str]:
|
57 |
+
"""Download README.md file for a given item."""
|
58 |
+
item_id = item['id']
|
59 |
+
raw_url = f"https://huggingface.co/{item_id}/raw/main/README.md"
|
60 |
+
try:
|
61 |
+
async with session.get(raw_url) as response:
|
62 |
+
if response.status == 200:
|
63 |
+
content = await response.text()
|
64 |
+
return item_id.replace('/', '_'), content
|
65 |
+
return item_id.replace('/', '_'), f"# Error downloading README for {item_id}\nStatus code: {response.status}"
|
66 |
+
except Exception as e:
|
67 |
+
return item_id.replace('/', '_'), f"# Error downloading README for {item_id}\nError: {str(e)}"
|
68 |
+
|
69 |
+
async def download_all_readmes(data: List[Dict]) -> tuple[str, str]:
|
70 |
+
"""Download all README files and create a zip archive."""
|
71 |
+
if not data:
|
72 |
+
return "", "No results to download"
|
73 |
+
|
74 |
+
zip_buffer = io.BytesIO()
|
75 |
+
status_message = "Downloading READMEs..."
|
76 |
+
|
77 |
+
async with aiohttp.ClientSession() as session:
|
78 |
+
tasks = [download_readme(session, item) for item in data]
|
79 |
+
results = await asyncio.gather(*tasks)
|
80 |
+
|
81 |
+
with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
|
82 |
+
for filename, content in results:
|
83 |
+
zip_file.writestr(f"{filename}.md", content)
|
84 |
+
|
85 |
+
zip_buffer.seek(0)
|
86 |
+
base64_zip = base64.b64encode(zip_buffer.getvalue()).decode()
|
87 |
+
|
88 |
+
download_link = f"""
|
89 |
+
<div style="margin-top: 10px;">
|
90 |
+
<a href="data:application/zip;base64,{base64_zip}"
|
91 |
+
download="readmes.zip"
|
92 |
+
style="display: inline-block; padding: 10px 20px;
|
93 |
+
background-color: #4CAF50; color: white;
|
94 |
+
text-decoration: none; border-radius: 5px;">
|
95 |
+
📥 Download READMEs Archive
|
96 |
+
</a>
|
97 |
+
</div>
|
98 |
+
"""
|
99 |
+
|
100 |
+
return download_link, "READMEs ready for download!"
|
101 |
+
|
102 |
+
def download_repository(repo_id: str, repo_type: str, temp_dir: str) -> str:
|
103 |
+
"""Download a single repository."""
|
104 |
+
try:
|
105 |
+
repo_path = snapshot_download(
|
106 |
+
repo_id=repo_id,
|
107 |
+
repo_type=repo_type.lower()[:-1], # Remove 's' from 'Models'/'Datasets'/'Spaces'
|
108 |
+
local_dir=os.path.join(temp_dir, repo_id.replace('/', '_')),
|
109 |
+
ignore_patterns=["*.bin", "*.pt", "*.pth", "*.ckpt", "*.safetensors"] # Ignore large binary files
|
110 |
+
)
|
111 |
+
return repo_path
|
112 |
+
except Exception as e:
|
113 |
+
print(f"Error downloading {repo_id}: {str(e)}")
|
114 |
+
return None
|
115 |
+
|
116 |
+
def create_repo_zip(data: List[Dict], search_type: str) -> tuple[str, str]:
|
117 |
+
"""Download repositories and create a zip archive."""
|
118 |
+
if not data:
|
119 |
+
return "", "No repositories to download"
|
120 |
+
|
121 |
+
# Create temporary directory
|
122 |
+
with tempfile.TemporaryDirectory() as temp_dir:
|
123 |
+
successful_downloads = []
|
124 |
+
|
125 |
+
# Download each repository
|
126 |
+
for item in data:
|
127 |
+
repo_path = download_repository(item['id'], search_type, temp_dir)
|
128 |
+
if repo_path:
|
129 |
+
successful_downloads.append(repo_path)
|
130 |
+
|
131 |
+
if not successful_downloads:
|
132 |
+
return "", "No repositories were successfully downloaded"
|
133 |
+
|
134 |
+
# Create zip file in memory
|
135 |
+
zip_buffer = io.BytesIO()
|
136 |
+
with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
|
137 |
+
for repo_path in successful_downloads:
|
138 |
+
repo_name = os.path.basename(repo_path)
|
139 |
+
for root, _, files in os.walk(repo_path):
|
140 |
+
for file in files:
|
141 |
+
file_path = os.path.join(root, file)
|
142 |
+
arcname = os.path.join(repo_name, os.path.relpath(file_path, repo_path))
|
143 |
+
zip_file.write(file_path, arcname)
|
144 |
+
|
145 |
+
# Convert to base64
|
146 |
+
zip_buffer.seek(0)
|
147 |
+
base64_zip = base64.b64encode(zip_buffer.getvalue()).decode()
|
148 |
+
|
149 |
+
download_link = f"""
|
150 |
+
<div style="margin-top: 10px;">
|
151 |
+
<a href="data:application/zip;base64,{base64_zip}"
|
152 |
+
download="repositories.zip"
|
153 |
+
style="display: inline-block; padding: 10px 20px;
|
154 |
+
background-color: #4CAF50; color: white;
|
155 |
+
text-decoration: none; border-radius: 5px;">
|
156 |
+
📥 Download Repositories Archive
|
157 |
+
</a>
|
158 |
+
</div>
|
159 |
+
"""
|
160 |
+
|
161 |
+
return download_link, f"Successfully downloaded {len(successful_downloads)} repositories"
|
162 |
+
|
163 |
+
def display_results(df: pd.DataFrame):
|
164 |
+
if df is not None and not df.empty:
|
165 |
+
html = "<div style='max-height: 400px; overflow-y: auto;'>"
|
166 |
+
for _, row in df.iterrows():
|
167 |
+
html += row['formatted_link']
|
168 |
+
html += "</div>"
|
169 |
+
return html
|
170 |
+
else:
|
171 |
+
return "<p>No results found.</p>"
|
172 |
+
|
173 |
+
def SwarmyTime(data: List[Dict]) -> Dict:
|
174 |
+
"""Aggregates all content from the given data."""
|
175 |
+
aggregated = {
|
176 |
+
"total_items": len(data),
|
177 |
+
"unique_authors": set(),
|
178 |
+
"total_downloads": 0,
|
179 |
+
"item_types": {"Models": 0, "Datasets": 0, "Spaces": 0}
|
180 |
+
}
|
181 |
+
|
182 |
+
for item in data:
|
183 |
+
aggregated["unique_authors"].add(item.get("author", "Unknown"))
|
184 |
+
aggregated["total_downloads"] += item.get("downloads", 0)
|
185 |
+
|
186 |
+
if "modelId" in item:
|
187 |
+
aggregated["item_types"]["Models"] += 1
|
188 |
+
elif "dataset" in item.get("id", ""):
|
189 |
+
aggregated["item_types"]["Datasets"] += 1
|
190 |
+
else:
|
191 |
+
aggregated["item_types"]["Spaces"] += 1
|
192 |
+
|
193 |
+
aggregated["unique_authors"] = len(aggregated["unique_authors"])
|
194 |
+
return aggregated
|
195 |
+
|
196 |
+
with gr.Blocks() as demo:
|
197 |
+
gr.Markdown("## Search the Hugging Face Hub")
|
198 |
+
with gr.Row():
|
199 |
+
search_query = gr.Textbox(label="Search Query", value="awacke1")
|
200 |
+
search_type = gr.Radio(["Models", "Datasets", "Spaces"], label="Search Type", value="Models")
|
201 |
+
search_button = gr.Button("Search")
|
202 |
+
|
203 |
+
results_html = gr.HTML(label="Search Results")
|
204 |
+
|
205 |
+
with gr.Row():
|
206 |
+
download_readme_button = gr.Button("📚 Download READMEs", visible=False)
|
207 |
+
download_repo_button = gr.Button("📦 Download Repositories", visible=False)
|
208 |
+
|
209 |
+
download_status = gr.Markdown("", label="Download Status")
|
210 |
+
download_area = gr.HTML("", label="Download Link")
|
211 |
+
metadata_output = gr.Textbox(label="Metadata", lines=10)
|
212 |
+
aggregated_output = gr.JSON(label="Aggregated Content")
|
213 |
+
|
214 |
+
search_type_state = gr.State("")
|
215 |
+
current_results = gr.State([])
|
216 |
+
|
217 |
+
def search_and_aggregate(query, search_type):
|
218 |
+
df = search_hub(query, search_type)
|
219 |
+
data = df.to_dict('records')
|
220 |
+
aggregated = SwarmyTime(data)
|
221 |
+
html_results = display_results(df)
|
222 |
+
show_download = len(data) > 0
|
223 |
+
return [
|
224 |
+
html_results, # results_html
|
225 |
+
show_download, # download_readme_button visible
|
226 |
+
show_download, # download_repo_button visible
|
227 |
+
"", # download_status
|
228 |
+
"", # download_area
|
229 |
+
aggregated, # aggregated_output
|
230 |
+
search_type, # search_type_state
|
231 |
+
data # current_results
|
232 |
+
]
|
233 |
+
|
234 |
+
async def handle_readme_download(data):
|
235 |
+
if not data:
|
236 |
+
return ["No results to download", ""]
|
237 |
+
download_link, status = await download_all_readmes(data)
|
238 |
+
return [status, download_link]
|
239 |
+
|
240 |
+
def handle_repo_download(data, search_type):
|
241 |
+
if not data:
|
242 |
+
return ["No results to download", ""]
|
243 |
+
download_link, status = create_repo_zip(data, search_type)
|
244 |
+
return [status, download_link]
|
245 |
+
|
246 |
+
search_button.click(
|
247 |
+
search_and_aggregate,
|
248 |
+
inputs=[search_query, search_type],
|
249 |
+
outputs=[
|
250 |
+
results_html,
|
251 |
+
download_readme_button,
|
252 |
+
download_repo_button,
|
253 |
+
download_status,
|
254 |
+
download_area,
|
255 |
+
aggregated_output,
|
256 |
+
search_type_state,
|
257 |
+
current_results
|
258 |
+
]
|
259 |
+
)
|
260 |
+
|
261 |
+
download_readme_button.click(
|
262 |
+
handle_readme_download,
|
263 |
+
inputs=[current_results],
|
264 |
+
outputs=[download_status, download_area]
|
265 |
+
)
|
266 |
+
|
267 |
+
download_repo_button.click(
|
268 |
+
handle_repo_download,
|
269 |
+
inputs=[current_results, search_type_state],
|
270 |
+
outputs=[download_status, download_area]
|
271 |
+
)
|
272 |
+
|
273 |
+
demo.launch(debug=True)
|