chore: onnx model

app.py

@@ -1,31 +1,128 @@
 import gradio as gr
+from optimum.pipelines import pipeline
 
 # モデルのロード
-model = gr.load("models/Miwa-Keita/zenz-v1-checkpoints")
+MODEL_NAME = "p1atdev/zenz-v1-onnx"
+pipe = pipeline("text-generation", MODEL_NAME)
+
+
+# ひらがなをカタカナに変換する関数
+def hiragana_to_katakana(hiragana: str):
+    katakana = ""
+    for char in hiragana:
+        # ひらがなの文字コードの範囲はU+3041からU+309F
+        if 0x3041 <= ord(char) <= 0x309F:
+            katakana += chr(ord(char) + 0x60)
+        else:
+            katakana += char
+    return katakana
+
 
 # 入力を調整する関数
-def preprocess_input(user_input):
+def preprocess_input(user_input: str):
     prefix = "\uEE00"  # 前に付与する文字列
     suffix = "\uEE01"  # 後ろに付与する文字列
-    processed_input = prefix + user_input + suffix
-    return model(processed_input)
+    processed_input = prefix + hiragana_to_katakana(user_input) + suffix
+
+    return processed_input
+
+
+# 出力を生成する関数
+def generate_output(input_text: str, num_beams: int = 4):
+    generated_outputs = pipe(
+        input_text,
+        max_new_tokens=len(input_text) * 2,
+        num_beams=num_beams,
+        num_return_sequences=num_beams,
+    )
+    generated_texts = [output["generated_text"] for output in generated_outputs]  # type: ignore
+    return generated_texts
+
 
 # 出力を調整する関数
-def postprocess_output(model_output):
+def postprocess_output(model_outputs: list[str]):
     suffix = "\uEE01"
     # \uEE01の後の部分を抽出
-    if suffix in model_output:
-        return model_output.split(suffix)[1]
-    return model_output
+    for i, model_output in enumerate(model_outputs):
+        if suffix in model_output:
+            model_outputs[i] = model_output.split(suffix)[1]
+    return "\n".join(
+        [f"{i+1}: {model_output}" for i, model_output in enumerate(model_outputs)]
+    )
+
+
+# 変換処理をまとめる関数
+def process_conversion(user_input: str, num_beams: int = 4):
+    processed_input = preprocess_input(user_input)
+    generated_outputs = generate_output(processed_input, num_beams)
+    postprocessed_output = postprocess_output(generated_outputs)
+    return postprocessed_output
+
 
 # インターフェースを定義
-iface = gr.Interface(
-    fn=lambda x: postprocess_output(preprocess_input(x)),
-    inputs=gr.Textbox(label="変換する文字列（カタカナ）"),
-    outputs=gr.Textbox(label="変換結果"),
-    title="ニューラルかな漢字変換モデルzenz-v1のデモ",
-    description="変換したい文字列をカタカナを入力してください"
-)
+def interface():
+    with gr.Blocks() as ui:
+        gr.Markdown(
+            """## ニューラルかな漢字変換モデルzenz-v1のデモ (ONNX版)
+            変換したい文字列をひらがな・カタカナを入力してください"""
+        )
+
+        with gr.Row():
+            with gr.Column():
+                input_text = gr.TextArea(
+                    label="変換する文字列（ひらがな・カタカナ）",
+                    info="変換したいテキストをひらがなかカタカナで入力します。入力すると右に反映されます。",
+                )
+                num_beams = gr.Slider(
+                    label="候補数",
+                    info="多くするとより変換に時間がかかります",
+                    minimum=1,
+                    maximum=20,
+                    step=1,
+                    value=4,
+                )
+
+            with gr.Column():
+                output_text = gr.TextArea(
+                    label="変換結果 (リアルタイム反映)",
+                    info="変換候補が出力されます。上の候補ほど確信度が高いです。",
+                )
+
+        gr.Examples(
+            examples=[
+                ["きめつのえいがをみました"],
+                ["はがいたいのでしかいにみてもらった"],
+                ["くつろぐにふといでかんたといいます"],
+                ["けんかをかった"],
+                ["けんかにかった"],
+                ["こうえんをおねがいする"],
+                ["こうえんでおねがいする"],
+                ["つきむらてまり"],
+            ],
+            inputs=[input_text],
+        )
+
+        gr.Markdown(
+            """\
+- 使用しているモデル (ONNX): [p1atdev/zenz-v1-onnx](https://huggingface.co/p1atdev/zenz-v1-onnx)
+- オリジナル(変換元)のモデル: [Miwa-Keita/zenz-v1-checkpoints](https://huggingface.co/Miwa-Keita/zenz-v1-checkpoints)
+"""
+        )
+
+        input_text.change(
+            fn=process_conversion,
+            inputs=[input_text, num_beams],
+            outputs=output_text,
+        )
+        num_beams.change(
+            fn=process_conversion,
+            inputs=[input_text, num_beams],
+            outputs=output_text,
+        )
+
+    ui.launch()
+
 
 # ローンチ
-iface.launch()
+if __name__ == "__main__":
+    interface()

conversion.py

@@ -0,0 +1,7 @@
+from optimum.onnxruntime import ORTModelForCausalLM
+
+MODEL_NAME = "Miwa-Keita/zenz-v1-checkpoints"
+
+ort_model = ORTModelForCausalLM.from_pretrained(MODEL_NAME, export=True)
+
+ort_model.save_pretrained(save_directory="./onnx")

requirements.txt

@@ -0,0 +1,2 @@
+torch
+optimum[onnxruntime]