app.py

import gradio as gr
import onnxruntime as ort
import numpy as np
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("cahya/distilbert-base-indonesian")
providers = ort.get_available_providers()
print("Available providers:", providers)

if "CUDAExecutionProvider" in providers:
    sess_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
else:
    sess_providers = ["CPUExecutionProvider"]

session = ort.InferenceSession("imphnenbert.onnx", providers=sess_providers)
print("Loaded ONNX session")

label = ["NORMAL", "PDF", "SH"]

def predict(text):
    encoded = tokenizer(text, padding=True, truncation=True, return_tensors="np")
    onnx_inputs = {}
    for input_meta in session.get_inputs():
        name = input_meta.name
        if name in encoded:
            onnx_inputs[name] = encoded[name]
        else:
            if "input_ids" in name and "input_ids" in encoded:
                onnx_inputs[name] = encoded["input_ids"]
            elif "attention" in name and "attention_mask" in encoded:
                onnx_inputs[name] = encoded["attention_mask"]
            elif "token_type" in name and "token_type_ids" in encoded:
                onnx_inputs[name] = encoded["token_type_ids"]
            else:
                raise KeyError(f"Couldn't map ONNX input '{name}' to tokenizer outputs. "
                               f"Available tokenizer keys: {list(encoded.keys())}")
            
    outputs = session.run(None, onnx_inputs)
    preds = np.argmax(outputs[0], axis=-1)

    return label[preds[0]]

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            text = gr.TextArea(label="Input Text", placeholder="Input Text")
            button = gr.Button("Predict")
        with gr.Column():
            output = gr.TextArea(label="Output")
    button.click(predict, [text], output)


demo.queue().launch()