|
import os |
|
import torch |
|
import gradio as gr |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
from huggingface_hub import HfApi, create_repo |
|
|
|
|
|
def convert_to_onnx_and_deploy(model_repo, input_text, hf_token, repo_name): |
|
try: |
|
|
|
os.environ["HUGGINGFACE_HUB_TOKEN"] = hf_token |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_repo) |
|
model = AutoModelForCausalLM.from_pretrained(model_repo) |
|
|
|
|
|
inputs = tokenizer(input_text, return_tensors="pt") |
|
|
|
|
|
onnx_file = f"{repo_name}.onnx" |
|
torch.onnx.export( |
|
model, |
|
inputs['input_ids'], |
|
onnx_file, |
|
input_names=['input_ids'], |
|
output_names=['output'], |
|
dynamic_axes={'input_ids': {0: 'batch_size'}, 'output': {0: 'batch_size'}} |
|
) |
|
|
|
|
|
api = HfApi() |
|
create_repo(repo_name, private=True) |
|
api.upload_file(onnx_file, repo_id=repo_name) |
|
|
|
return f"ONNXモデルが作成され、リポジトリにデプロイされました: {repo_name}" |
|
except Exception as e: |
|
return str(e) |
|
|
|
|
|
iface = gr.Interface( |
|
fn=convert_to_onnx_and_deploy, |
|
inputs=[ |
|
gr.Textbox(label="モデルリポジトリ(例: rinna/japanese-gpt2-medium)"), |
|
gr.Textbox(label="入力テキスト"), |
|
gr.Textbox(label="Hugging Faceトークン", type="password"), |
|
gr.Textbox(label="デプロイ先リポジトリ名") |
|
], |
|
outputs="text", |
|
title="ONNX変換とモデルデプロイ機能", |
|
description="指定したHugging FaceのモデルリポジトリをONNX形式に変換し、デプロイします。" |
|
) |
|
|
|
|
|
iface.launch(server_port=7865) |