chore: remote flag of model backend

app.py

@@ -5,7 +5,7 @@ import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# from optimum.onnxruntime import ORTModelForCausalLM
+from optimum.onnxruntime import ORTModelForCausalLM
 
 from huggingface_hub import login
 
@@ -17,8 +17,14 @@ MODEL_NAME = (
     else "p1atdev/dart-v1-sft"
 )
 HF_READ_TOKEN = os.environ.get("HF_READ_TOKEN")
+MODEL_BACKEND = (
+    os.environ.get("MODEL_BACKEND")
+    if os.environ.get("MODEL_BACKEND") is not None
+    else "ONNX (quantized)"
+)
 
 assert isinstance(MODEL_NAME, str)
+assert isinstance(MODEL_BACKEND, str)
 
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_NAME,
@@ -30,19 +36,19 @@ model = {
         MODEL_NAME,
         token=HF_READ_TOKEN,
     ),
-    # "ort": ORTModelForCausalLM.from_pretrained(
-    #     MODEL_NAME,
-    # ),
-    # "ort_qantized": ORTModelForCausalLM.from_pretrained(
-    #     MODEL_NAME,
-    #     file_name="model_quantized.onnx",
-    # ),
+    "ort": ORTModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+    ),
+    "ort_qantized": ORTModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        file_name="model_quantized.onnx",
+    ),
 }
 
 MODEL_BACKEND_MAP = {
     "Default": "default",
-    # "ONNX (normal)": "ort",
-    # "ONNX (quantized)": "ort_qantized",
+    "ONNX (normal)": "ort",
+    "ONNX (quantized)": "ort_qantized",
 }
 
 try:
@@ -288,7 +294,7 @@ def handle_inputs(
     top_p: float = 1.0,
     top_k: int = 20,
     num_beams: int = 1,
-    model_backend: str = "Default",
+    # model_backend: str = "Default",
 ):
     """
     Returns:
@@ -340,7 +346,7 @@ def handle_inputs(
 
     generated_ids = generate(
         prompt,
-        model_backend=model_backend,
+        model_backend=MODEL_BACKEND,
         max_new_tokens=max_new_tokens,
         min_new_tokens=min_new_tokens,
         do_sample=True,
@@ -395,21 +401,30 @@ def demo():
     with gr.Blocks() as ui:
         gr.Markdown(
             """\
-# Danbooru Tags Transformer Demo """
+# Danbooru Tags Transformer Demo 
+
+Collection: [Dart (Danbooru Tags Transformer)](https://huggingface.co/collections/p1atdev/dart-danbooru-tags-transformer-65d687604ff57dc62ae40945)
+
+Models:
+
+- [p1atdev/dart-v1-sft](https://huggingface.co/p1atdev/dart-v1-sft)
+- [p1atdev/dart-v1-base](https://huggingface.co/p1atdev/dart-v1-base)
+
+"""
         )
 
         with gr.Row():
             with gr.Column():
 
-                with gr.Group(
-                    visible=False,
-                ):
-                    model_backend_radio = gr.Radio(
-                        label="Model backend",
-                        choices=list(MODEL_BACKEND_MAP.keys()),
-                        value="Default",
-                        interactive=True,
-                    )
+                # with gr.Group(
+                #     visible=False,
+                # ):
+                #     model_backend_radio = gr.Radio(
+                #         label="Model backend",
+                #         choices=list(MODEL_BACKEND_MAP.keys()),
+                #         value="Default",
+                #         interactive=True,
+                #     )
 
                 with gr.Group():
                     rating_dropdown = gr.Dropdown(
@@ -663,7 +678,7 @@ def demo():
                     top_p_slider,
                     top_k_slider,
                     num_beams_slider,
-                    model_backend_radio,
+                    # model_backend_radio,
                 ],
                 outputs=[
                     output_tags_natural,

requirements.txt

@@ -1,2 +1,4 @@
 torch==2.1.0
-transformers==4.38.0
+accelerate==0.26.1
+transformers==4.38.0
+optimum[onnxruntime]==1.17.1