Spaces:
Running
Running
| import csv | |
| import os | |
| from datetime import datetime | |
| from typing import Optional, Union | |
| import gradio as gr | |
| from huggingface_hub import HfApi, Repository | |
| from optimum_neuron_export import convert | |
| from gradio_huggingfacehub_search import HuggingfaceHubSearch | |
| from apscheduler.schedulers.background import BackgroundScheduler | |
| DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/neuron-exports" | |
| DATA_FILENAME = "exports.csv" | |
| DATA_FILE = os.path.join("data", DATA_FILENAME) | |
| HF_TOKEN = os.environ.get("HF_WRITE_TOKEN") | |
| DATADIR = "neuron_exports_data" | |
| repo: Optional[Repository] = None | |
| # Uncomment if you want to push to dataset repo with token | |
| # if HF_TOKEN: | |
| # repo = Repository(local_dir=DATADIR, clone_from=DATASET_REPO_URL, token=HF_TOKEN) | |
| # Define all possible tasks and their categories for coloring | |
| TASK_CATEGORIES = { | |
| "auto": {"color": "#6b7280", "category": "Auto"}, | |
| "feature-extraction": {"color": "#3b82f6", "category": "Feature Extraction"}, | |
| "fill-mask": {"color": "#8b5cf6", "category": "NLP"}, | |
| "multiple-choice": {"color": "#8b5cf6", "category": "NLP"}, | |
| "question-answering": {"color": "#8b5cf6", "category": "NLP"}, | |
| "text-classification": {"color": "#8b5cf6", "category": "NLP"}, | |
| "token-classification": {"color": "#8b5cf6", "category": "NLP"}, | |
| "text-generation": {"color": "#10b981", "category": "Text Generation"}, | |
| "text2text-generation": {"color": "#10b981", "category": "Text Generation"}, | |
| "audio-classification": {"color": "#f59e0b", "category": "Audio"}, | |
| "automatic-speech-recognition": {"color": "#f59e0b", "category": "Audio"}, | |
| "audio-frame-classification": {"color": "#f59e0b", "category": "Audio"}, | |
| "audio-xvector": {"color": "#f59e0b", "category": "Audio"}, | |
| "image-classification": {"color": "#ef4444", "category": "Vision"}, | |
| "object-detection": {"color": "#ef4444", "category": "Vision"}, | |
| "semantic-segmentation": {"color": "#ef4444", "category": "Vision"}, | |
| "text-to-image": {"color": "#ec4899", "category": "Multimodal"}, | |
| "image-to-image": {"color": "#ec4899", "category": "Multimodal"}, | |
| "inpaint": {"color": "#ec4899", "category": "Multimodal"}, | |
| "zero-shot-image-classification": {"color": "#ec4899", "category": "Multimodal"}, | |
| "sentence-similarity": {"color": "#06b6d4", "category": "Similarity"}, | |
| "Auto": {"color": "#6b7280", "category": "Auto"}, | |
| "Feature Extraction": {"color": "#3b82f6", "category": "Feature Extraction"}, | |
| "NLP": {"color": "#8b5cf6", "category": "NLP"}, | |
| "Text Generation": {"color": "#10b981", "category": "Text Generation"}, | |
| "Audio": {"color": "#f59e0b", "category": "Audio"}, | |
| "Vision": {"color": "#ef4444", "category": "Vision"}, | |
| "Multimodal": {"color": "#ec4899", "category": "Multimodal"}, | |
| "Similarity": {"color": "#06b6d4", "category": "Similarity"}, | |
| } | |
| # Get all tasks for dropdown | |
| ALL_TASKS = list(TASK_CATEGORIES.keys()) | |
| def create_task_tag(task: str) -> str: | |
| """Create a colored HTML tag for a task""" | |
| if task in TASK_CATEGORIES: | |
| color = TASK_CATEGORIES[task]["color"] | |
| return f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>' | |
| else: | |
| return f'<span style="background-color: #6b7280; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>' | |
| def format_tasks_for_table(tasks_str: str) -> str: | |
| """Convert comma-separated tasks into colored tags""" | |
| tasks = [task.strip() for task in tasks_str.split(',')] | |
| return ' '.join([create_task_tag(task) for task in tasks]) | |
| def neuron_export(model_id: str, task: str, oauth_token: gr.OAuthToken) -> str: | |
| if oauth_token.token is None: | |
| return "You must be logged in to use this space" | |
| if not model_id: | |
| return f"### Invalid input 🐞 Please specify a model name, got {model_id}" | |
| try: | |
| api = HfApi(token=oauth_token.token) | |
| error, commit_info = convert(api=api, model_id=model_id, task=task, token=oauth_token.token) | |
| if error != "0": | |
| return error | |
| print("[commit_info]", commit_info) | |
| # Save in a private dataset if repo initialized | |
| if repo is not None: | |
| repo.git_pull(rebase=True) | |
| with open(os.path.join(DATADIR, DATA_FILE), "a") as csvfile: | |
| writer = csv.DictWriter( | |
| csvfile, fieldnames=["model_id", "pr_url", "time"] | |
| ) | |
| writer.writerow( | |
| { | |
| "model_id": model_id, | |
| "pr_url": commit_info.pr_url, | |
| "time": str(datetime.now()), | |
| } | |
| ) | |
| commit_url = repo.push_to_hub() | |
| print("[dataset]", commit_url) | |
| pr_revision = commit_info.pr_revision.replace("/", "%2F") | |
| return f"#### Success 🔥 Yay! This model was successfully exported and a PR was opened using your token: [{commit_info.pr_url}]({commit_info.pr_url}). If you would like to use the exported model without waiting for the PR to be approved, head to https://huggingface.co/{model_id}/tree/{pr_revision}" | |
| except Exception as e: | |
| return f"#### Error: {e}" | |
| TITLE_IMAGE = """ | |
| <div style="display: block; margin-left: auto; margin-right: auto; width: 50%;"> | |
| <img src="https://huggingface.co/spaces/optimum/neuron-export/resolve/main/huggingfaceXneuron.png"/> | |
| </div> | |
| """ | |
| TITLE = """ | |
| <div style="text-align: center; max-width: 1400px; margin: 0 auto;"> | |
| <h1 style="font-weight: 900; margin-bottom: 10px; margin-top: 10px; font-size: 2.2rem;"> | |
| 🤗 Optimum Neuron Model Exporter 🏎️ (WIP) | |
| </h1> | |
| </div> | |
| """ | |
| DESCRIPTION = """ | |
| This Space allows you to automatically export 🤗 transformers models hosted on the Hugging Face Hub to AWS Neuron-optimized format for Inferentia/Trainium acceleration. It opens a PR on the target model, and it is up to the owner of the original model to merge the PR to allow people to leverage Neuron optimization! | |
| **Features:** | |
| - Automatically opens PR with Neuron-optimized model | |
| - Preserves original model weights | |
| - Adds proper tags to model card | |
| **Requirements:** | |
| - Model must be compatible with [Optimum Neuron](https://huggingface.co/docs/optimum-neuron) | |
| - User must be logged in with write token | |
| """ | |
| # Custom CSS to fix dark mode compatibility and transparency issues | |
| CUSTOM_CSS = """ | |
| /* Fix for HuggingfaceHubSearch component visibility in both light and dark modes */ | |
| .gradio-container .gr-form { | |
| background: var(--background-fill-primary) !important; | |
| border: 1px solid var(--border-color-primary) !important; | |
| } | |
| /* Ensure text is visible in both modes */ | |
| .gradio-container input[type="text"], | |
| .gradio-container textarea, | |
| .gradio-container .gr-textbox input { | |
| color: var(--body-text-color) !important; | |
| background: var(--input-background-fill) !important; | |
| border: 1px solid var(--border-color-primary) !important; | |
| } | |
| /* Fix dropdown/search results visibility */ | |
| .gradio-container .gr-dropdown, | |
| .gradio-container .gr-dropdown .gr-box, | |
| .gradio-container [data-testid="textbox"] { | |
| background: var(--background-fill-primary) !important; | |
| color: var(--body-text-color) !important; | |
| border: 1px solid var(--border-color-primary) !important; | |
| } | |
| /* Fix for search component specifically */ | |
| .gradio-container .gr-form > div, | |
| .gradio-container .gr-form input { | |
| background: var(--input-background-fill) !important; | |
| color: var(--body-text-color) !important; | |
| } | |
| /* Ensure proper contrast for placeholder text */ | |
| .gradio-container input::placeholder { | |
| color: var(--body-text-color-subdued) !important; | |
| opacity: 0.7; | |
| } | |
| /* Fix any remaining transparent backgrounds */ | |
| .gradio-container .gr-box, | |
| .gradio-container .gr-panel { | |
| background: var(--background-fill-primary) !important; | |
| } | |
| /* Make sure search results are visible */ | |
| .gradio-container .gr-dropdown-item { | |
| color: var(--body-text-color) !important; | |
| background: var(--background-fill-primary) !important; | |
| } | |
| .gradio-container .gr-dropdown-item:hover { | |
| background: var(--background-fill-secondary) !important; | |
| } | |
| /* Task tag styling improvements */ | |
| .task-tags { | |
| line-height: 1.8; | |
| } | |
| .task-tags span { | |
| display: inline-block; | |
| margin: 2px; | |
| } | |
| """ | |
| with gr.Blocks(css=CUSTOM_CSS) as demo: | |
| # Login requirement notice and button | |
| gr.Markdown("**You must be logged in to use this space**") | |
| gr.LoginButton(min_width=250) | |
| # Centered title and image | |
| gr.HTML(TITLE_IMAGE) | |
| gr.HTML(TITLE) | |
| # Full-width description | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Tabs(): | |
| with gr.Tab("Export Model"): | |
| # Input controls in a row | |
| with gr.Row(): | |
| input_model = HuggingfaceHubSearch( | |
| label="Hub model ID", | |
| placeholder="Search for model ID on the hub", | |
| search_type="model", | |
| ) | |
| input_task = gr.Dropdown( | |
| choices=ALL_TASKS, | |
| value="auto", | |
| label='Task', | |
| info='Select the task type for your model (auto will infer automatically)', | |
| ) | |
| # Export button below the inputs | |
| btn = gr.Button("Export to Neuron", size="lg") | |
| # Output section | |
| output = gr.Markdown(label="Output") | |
| btn.click( | |
| fn=neuron_export, | |
| inputs=[input_model, input_task], | |
| outputs=output, | |
| ) | |
| with gr.Tab("Supported Architectures"): | |
| gr.HTML(f""" | |
| <div style="margin-bottom: 20px;"> | |
| <h3>🎨 Task Categories Legend</h3> | |
| <div class="task-tags"> | |
| {create_task_tag("Auto")} | |
| {create_task_tag("Feature Extraction")} | |
| {create_task_tag("NLP")} | |
| {create_task_tag("Text Generation")} | |
| {create_task_tag("Audio")} | |
| {create_task_tag("Vision")} | |
| {create_task_tag("Multimodal")} | |
| {create_task_tag("Similarity")} | |
| </div> | |
| </div> | |
| """) | |
| gr.HTML(f""" | |
| <h2>🤗 Transformers</h2> | |
| <table style="width: 100%; border-collapse: collapse; margin: 20px 0;"> | |
| <thead> | |
| <tr style="background-color: var(--background-fill-secondary);"> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ALBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">AST</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, audio-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">BERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">BLOOM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Beit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CamemBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CLIP</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvNext</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvNextV2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CvT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DeBERTa (INF2 only)</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DeBERTa-v2 (INF2 only)</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Deit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DistilBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DonutSwin</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Dpt</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ELECTRA</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ESM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">FlauBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">GPT2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Hubert</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Levit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Llama, Llama 2, Llama 3</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Mistral</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Mixtral</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileNetV2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification, semantic-segmentation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileViT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification, semantic-segmentation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ModernBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MPNet</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">OPT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Phi</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">RoBERTa</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">RoFormer</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Swin</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">T5</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text2text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">UniSpeech</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">UniSpeech-SAT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ViT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Wav2Vec2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">WavLM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Whisper</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("automatic-speech-recognition")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">XLM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">XLM-RoBERTa</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Yolos</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, object-detection")}</td></tr> | |
| </tbody> | |
| </table> | |
| <h2>🧨 Diffusers</h2> | |
| <table style="width: 100%; border-collapse: collapse; margin: 20px 0;"> | |
| <thead> | |
| <tr style="background-color: var(--background-fill-secondary);"> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion XL Base</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion XL Refiner</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">SDXL Turbo</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">LCM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">PixArt-α</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">PixArt-Σ</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr> | |
| </tbody> | |
| </table> | |
| <h2>🤖 Sentence Transformers</h2> | |
| <table style="width: 100%; border-collapse: collapse; margin: 20px 0;"> | |
| <thead> | |
| <tr style="background-color: var(--background-fill-secondary);"> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Transformer</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, sentence-similarity")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CLIP</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, zero-shot-image-classification")}</td></tr> | |
| </tbody> | |
| </table> | |
| <div style="margin-top: 20px;"> | |
| <p>💡 <strong>Note</strong>: Some architectures may have specific requirements or limitations. DeBERTa models are only supported on INF2 instances.</p> | |
| <p>For more details, check the <a href="https://huggingface.co/docs/optimum-neuron" target="_blank">Optimum Neuron documentation</a>.</p> | |
| </div> | |
| """) | |
| # Add spacing between tabs and content | |
| gr.Markdown("<br><br><br><br>") | |
| if __name__ == "__main__": | |
| def restart_space(): | |
| if HF_TOKEN: | |
| HfApi().restart_space(repo_id="optimum/neuron-export", token=HF_TOKEN, factory_reboot=True) | |
| scheduler = BackgroundScheduler() | |
| scheduler.add_job(restart_space, "interval", seconds=21600) # Restart every 6 hours | |
| scheduler.start() | |
| demo.launch() |