Spaces:
Running
Running
File size: 4,280 Bytes
11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 11e4904 58e6c10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
import csv
import os
from datetime import datetime
from typing import Optional
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)
def neuron_export(model_id: str, task: str) -> str:
if not model_id:
return f"### Invalid input 🐞 Please specify a model name, got {model_id}"
try:
api = HfApi(token=HF_TOKEN) # Use HF_TOKEN if available, else anonymous
token = HF_TOKEN # Pass token to convert only if available
error, commit_info = convert(api=api, model_id=model_id, task=task, 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 🔥 This model was successfully exported and a PR was opened: [{commit_info.pr_url}]({commit_info.pr_url}). To use the model before the PR is approved, go 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="display: inline-flex; align-items: center; text-align: center; max-width: 1400px; gap: 0.8rem; font-size: 2.2rem;">
<h1 style="font-weight: 900; margin-bottom: 10px; margin-top: 10px;">
🤗 Optimum Neuron Model Exporter
</h1>
</div>
"""
DESCRIPTION = """
Export 🤗 Transformers models hosted on the Hugging Face Hub to AWS Neuron-optimized format for Inferentia/Trainium acceleration.
*Features:*
- Automatically opens PR with Neuron-optimized model
- Preserves original model weights
- Adds proper tags to model card
*Note:*
- PR creation requires the Space owner to have a valid write token set via HF_WRITE_TOKEN
"""
with gr.Blocks() as demo:
gr.HTML(TITLE_IMAGE)
gr.HTML(TITLE)
with gr.Row():
with gr.Column(scale=50):
gr.Markdown(DESCRIPTION)
with gr.Column(scale=50):
input_model = HuggingfaceHubSearch(
label="Hub model ID",
placeholder="Search for model ID on the hub",
search_type="model",
)
input_task = gr.Textbox(
value="auto",
max_lines=1,
label='Task (can be left to "auto", will be automatically inferred)',
)
btn = gr.Button("Export to Neuron")
output = gr.Markdown(label="Output")
btn.click(
fn=neuron_export,
inputs=[input_model, input_task],
outputs=output,
)
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)
scheduler.start()
demo.launch()
|