awacke1 commited on
Commit
e5c4bb6
·
verified ·
1 Parent(s): 7ab7b3a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +273 -0
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)