Spaces:
Runtime error
Runtime error
| import os | |
| import gradio as gr | |
| import numpy as np | |
| import onnxruntime as ort | |
| import sentencepiece as spm | |
| from huggingface_hub import snapshot_download | |
| # ========================= | |
| # DOWNLOAD MODEL | |
| # ========================= | |
| MODEL_ID = "Phitran21/AnViMt" | |
| print("Đang tải AnViMt...") | |
| MODEL_DIR = snapshot_download( | |
| repo_id=MODEL_ID, | |
| repo_type="model", | |
| ) | |
| print("Model directory:", MODEL_DIR) | |
| # ========================= | |
| # PATH | |
| # ========================= | |
| ENCODER_PATH = os.path.join( | |
| MODEL_DIR, | |
| "encoder_model.onnx" | |
| ) | |
| DECODER_PATH = os.path.join( | |
| MODEL_DIR, | |
| "decoder_model.onnx" | |
| ) | |
| SPM_PATH = os.path.join( | |
| MODEL_DIR, | |
| "spm.model" | |
| ) | |
| # ========================= | |
| # SETTINGS | |
| # ========================= | |
| MAX_INPUT_LEN = 256 | |
| MAX_OUTPUT_LEN = 256 | |
| PAD_ID = 0 | |
| UNK_ID = 1 | |
| BOS_ID = 2 | |
| EOS_ID = 3 | |
| # ========================= | |
| # CHECK FILES | |
| # ========================= | |
| for path in [ | |
| ENCODER_PATH, | |
| DECODER_PATH, | |
| SPM_PATH, | |
| ]: | |
| if not os.path.exists(path): | |
| raise FileNotFoundError( | |
| f"Không tìm thấy file: {path}" | |
| ) | |
| # ========================= | |
| # LOAD TOKENIZER | |
| # ========================= | |
| print("Đang load tokenizer...") | |
| sp = spm.SentencePieceProcessor() | |
| sp.load(SPM_PATH) | |
| print( | |
| f"Tokenizer loaded | " | |
| f"vocab={sp.get_piece_size()} | " | |
| f"PAD={sp.pad_id()} " | |
| f"BOS={sp.bos_id()} " | |
| f"EOS={sp.eos_id()}" | |
| ) | |
| # ========================= | |
| # LOAD ONNX MODEL | |
| # ========================= | |
| print("Đang load ONNX model...") | |
| providers = [ | |
| "CPUExecutionProvider" | |
| ] | |
| encoder_session = ort.InferenceSession( | |
| ENCODER_PATH, | |
| providers=providers | |
| ) | |
| decoder_session = ort.InferenceSession( | |
| DECODER_PATH, | |
| providers=providers | |
| ) | |
| print("AnViMt ready!") | |
| # ========================= | |
| # TRANSLATE | |
| # ========================= | |
| def translate(text): | |
| text = text.strip() | |
| if not text: | |
| return "" | |
| # ------------------------- | |
| # ENCODE INPUT | |
| # ------------------------- | |
| src_ids = sp.encode( | |
| text, | |
| out_type=int | |
| ) | |
| src_ids = src_ids[ | |
| :MAX_INPUT_LEN - 2 | |
| ] | |
| src_ids = ( | |
| [BOS_ID] | |
| + src_ids | |
| + [EOS_ID] | |
| ) | |
| input_ids = np.array( | |
| [src_ids], | |
| dtype=np.int64 | |
| ) | |
| attention_mask = np.ones_like( | |
| input_ids, | |
| dtype=np.int64 | |
| ) | |
| # ------------------------- | |
| # ENCODER | |
| # ------------------------- | |
| encoder_outputs = encoder_session.run( | |
| None, | |
| { | |
| "input_ids": input_ids, | |
| "attention_mask": attention_mask, | |
| } | |
| ) | |
| encoder_hidden_states = encoder_outputs[0] | |
| # ------------------------- | |
| # DECODER GREEDY | |
| # ------------------------- | |
| generated = [BOS_ID] | |
| for _ in range(MAX_OUTPUT_LEN): | |
| decoder_input_ids = np.array( | |
| [generated], | |
| dtype=np.int64 | |
| ) | |
| decoder_inputs = {} | |
| for inp in decoder_session.get_inputs(): | |
| name = inp.name | |
| if name == "input_ids": | |
| decoder_inputs[name] = ( | |
| decoder_input_ids | |
| ) | |
| elif ( | |
| "encoder_hidden_states" | |
| in name | |
| ): | |
| decoder_inputs[name] = ( | |
| encoder_hidden_states | |
| ) | |
| elif ( | |
| "encoder_attention_mask" | |
| in name | |
| ): | |
| decoder_inputs[name] = ( | |
| attention_mask | |
| ) | |
| decoder_outputs = decoder_session.run( | |
| None, | |
| decoder_inputs | |
| ) | |
| logits = decoder_outputs[0] | |
| next_token_logits = ( | |
| logits[0, -1, :] | |
| ) | |
| next_token = int( | |
| np.argmax( | |
| next_token_logits | |
| ) | |
| ) | |
| generated.append( | |
| next_token | |
| ) | |
| if next_token == EOS_ID: | |
| break | |
| # ------------------------- | |
| # DECODE OUTPUT | |
| # ------------------------- | |
| output_ids = generated[1:] | |
| if EOS_ID in output_ids: | |
| output_ids = output_ids[ | |
| :output_ids.index(EOS_ID) | |
| ] | |
| result = sp.decode( | |
| output_ids | |
| ) | |
| return result | |
| # ========================= | |
| # GRADIO UI | |
| # ========================= | |
| with gr.Blocks( | |
| title="AnViMt - English to Vietnamese Translator" | |
| ) as demo: | |
| gr.Markdown( | |
| """ | |
| # 🇻🇳 AnViMt | |
| ### English → Vietnamese Machine Translation | |
| Nhập văn bản tiếng Anh và AnViMt sẽ dịch sang tiếng Việt. | |
| **Model:** Phitran21/AnViMt | |
| **Architecture:** BART Seq2Seq | |
| **Inference:** ONNX Runtime | |
| """ | |
| ) | |
| with gr.Row(): | |
| input_text = gr.Textbox( | |
| label="English", | |
| placeholder=( | |
| "Enter English text..." | |
| ), | |
| lines=8 | |
| ) | |
| output_text = gr.Textbox( | |
| label="Vietnamese", | |
| lines=8 | |
| ) | |
| translate_button = gr.Button( | |
| "Translate", | |
| variant="primary" | |
| ) | |
| translate_button.click( | |
| fn=translate, | |
| inputs=input_text, | |
| outputs=output_text | |
| ) | |
| input_text.submit( | |
| fn=translate, | |
| inputs=input_text, | |
| outputs=output_text | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| [ | |
| "The system automatically restarts the service after an unexpected process crash." | |
| ], | |
| [ | |
| "A lone swordsman entered the abandoned temple at midnight, unaware that his greatest enemy was waiting in the shadows." | |
| ], | |
| [ | |
| "Could you please restart the service and check whether the error still occurs?" | |
| ], | |
| ], | |
| inputs=input_text | |
| ) | |
| demo.launch() |