Create onnx_infer.py

onnx_infer.py

@@ -0,0 +1,115 @@
+import time
+import torch 
+import numpy as np
+import onnxruntime as ort
+from transformers import AutoTokenizer
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+tokenizer = AutoTokenizer.from_pretrained("b3x0m/xomdich-tokenizer")
+
+onnx_model_path = "hyper-xomdich.onnx"
+so = ort.SessionOptions()
+so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+so.enable_mem_pattern = True
+so.enable_cpu_mem_arena = True
+session = ort.InferenceSession(onnx_model_path, so, providers=['CPUExecutionProvider', 'CUDAExecutionProvider'])
+
+def prepare_input(text, tokenizer, max_length=512):
+    text = text + '。' if '。' not in text else text
+    encoder_inputs = tokenizer(
+        f"{text}</s>",
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+        add_special_tokens=True
+    )
+
+    decoder_input = torch.tensor([[tokenizer.bos_token_id]])
+    src_len = encoder_inputs["input_ids"].size(1)
+    tgt_len = decoder_input.size(1)
+    
+    encoder_attention_mask = encoder_inputs["attention_mask"].unsqueeze(1).unsqueeze(1)
+    encoder_attention_mask = encoder_attention_mask.expand(-1, -1, src_len, src_len)
+    encoder_attention_mask = (1.0 - encoder_attention_mask) * -10000.0
+    
+    decoder_attention_mask = torch.triu(torch.ones((tgt_len, tgt_len)), diagonal=1).bool()
+    decoder_attention_mask = decoder_attention_mask.unsqueeze(0).unsqueeze(0)
+    decoder_attention_mask = decoder_attention_mask.float() * -10000.0
+    
+    cross_attention_mask = encoder_inputs["attention_mask"].unsqueeze(1).unsqueeze(1)
+    cross_attention_mask = cross_attention_mask.expand(-1, -1, tgt_len, src_len)
+    cross_attention_mask = (1.0 - cross_attention_mask) * -10000.0
+    
+    inputs = {
+        "input_ids": encoder_inputs["input_ids"].numpy().astype(np.int64),
+        "attention_mask": encoder_attention_mask.numpy().astype(np.float32),
+        "decoder_input_ids": decoder_input.numpy().astype(np.int64),
+        "decoder_attention_mask": decoder_attention_mask.numpy().astype(np.float32),
+        "cross_attention_mask": cross_attention_mask.numpy().astype(np.float32)
+    }
+    return inputs, decoder_input, src_len
+
+def stable_softmax(logits):
+    logits = logits - np.max(logits, axis=-1, keepdims=True)
+    exp_logits = np.exp(logits)
+    probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
+    return probs
+
+def generate_translation_onnx(text, session=session, tokenizer=tokenizer, max_length=512, temperature=0.1):
+    inputs, decoder_input, src_len = prepare_input(text, tokenizer, max_length)
+    generated_tokens = [tokenizer.bos_token_id]
+    
+    start_time = time.time()
+    for step in range(max_length):
+        outputs = session.run(["logits"], inputs)
+        logits = outputs[0]
+
+        # temperature sampling
+        next_token_logits = logits[:, -1, :] / temperature
+        probs = stable_softmax(next_token_logits)
+        next_token = np.random.choice(len(probs[0]), p=probs[0])
+        
+        # greedy search
+        # next_token = logits[:, -1, :].argmax(-1)
+        # next_token = next_token.item()
+
+        generated_tokens.append(next_token)
+        
+        if next_token == tokenizer.eos_token_id:
+            break
+        
+        decoder_input = torch.tensor([generated_tokens])
+        tgt_len = decoder_input.size(1)
+        
+        decoder_attention_mask = torch.triu(torch.ones((tgt_len, tgt_len)), diagonal=1).bool()
+        decoder_attention_mask = decoder_attention_mask.unsqueeze(0).unsqueeze(0)
+        decoder_attention_mask = decoder_attention_mask.float() * -10000.0
+        
+        cross_attention_mask = torch.ones((1, 1, tgt_len, src_len))
+        cross_attention_mask = (1.0 - cross_attention_mask) * -10000.0
+        
+        inputs["decoder_input_ids"] = decoder_input.numpy().astype(np.int64)
+        inputs["decoder_attention_mask"] = decoder_attention_mask.numpy().astype(np.float32)
+        inputs["cross_attention_mask"] = cross_attention_mask.numpy().astype(np.float32)
+
+    duration = time.time() - start_time
+    speed = len(text) / duration if duration > 0 else 0
+
+    performance_info = f"Time: {duration:.2f}s | Speed: {speed:.2f} chars/s"
+    
+    return tokenizer.decode(generated_tokens, skip_special_tokens=True), performance_info
+
+def interactive_translation():
+    while True:
+        text = input("Text input (press q to exit): ")
+
+        if text.lower() == 'q':
+            break
+        
+        translation, performance = generate_translation_onnx(text, session, tokenizer)
+        print(f"Translation: {translation}")
+        print(f"{performance}")
+
+if __name__ == "__main__":
+    interactive_translation()