File size: 5,507 Bytes
4966a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | import os
import sys
import json
import time
# Ensure user packages are in path for onnxruntime
sys.path.append("/home/hypr4/.local/lib/python3.12/site-packages")
import numpy as np
import onnxruntime as ort
class ONNXTransliterator:
def __init__(self, model_dir):
# Load vocabularies
with open(os.path.join(model_dir, "input_vocab.json"), "r", encoding="utf-8") as f:
self.src_vocab = json.load(f)
with open(os.path.join(model_dir, "target_vocab.json"), "r", encoding="utf-8") as f:
self.tgt_vocab = json.load(f)
self.src_idx2char = {v: k for k, v in self.src_vocab.items()}
self.tgt_idx2char = {v: k for k, v in self.tgt_vocab.items()}
# Load ONNX (force 1-thread CPU mode)
opts = ort.SessionOptions()
opts.intra_op_num_threads = 1
opts.inter_op_num_threads = 1
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
self.encoder_sess = ort.InferenceSession(os.path.join(model_dir, "encoder.onnx"), sess_options=opts)
self.decoder_sess = ort.InferenceSession(os.path.join(model_dir, "decoder.onnx"), sess_options=opts)
def transliterate_word(self, word):
src_ids = [self.src_vocab["<s>"]]
for char in word:
src_ids.append(self.src_vocab.get(char, self.src_vocab["<unk>"]))
src_ids.append(self.src_vocab["</s>"])
input_ids = np.array([src_ids], dtype=np.int64)
enc_outputs, enc_h, enc_c = self.encoder_sess.run(
["encoder_outputs", "h_states", "c_states"],
{"input_ids": input_ids}
)
num_layers = 2
hidden_dim = 256
dec_h = np.zeros((num_layers, 1, hidden_dim), dtype=np.float32)
dec_c = np.zeros((num_layers, 1, hidden_dim), dtype=np.float32)
for i in range(num_layers):
dec_h[i] = (enc_h[2*i] + enc_h[2*i+1]) / 2.0
dec_c[i] = (enc_c[2*i] + enc_c[2*i+1]) / 2.0
dec_input = np.array([self.tgt_vocab["<s>"]], dtype=np.int64)
output_chars = []
for step in range(32):
logits, dec_h, dec_c, _ = self.decoder_sess.run(
["logits", "h", "c", "attn_weights"],
{
"input_char": dec_input,
"prev_h": dec_h,
"prev_c": dec_c,
"encoder_outputs": enc_outputs
}
)
next_char_idx = int(np.argmax(logits[0]))
if next_char_idx == self.tgt_vocab["</s>"] or next_char_idx == self.tgt_vocab["<pad>"]:
break
output_chars.append(self.tgt_idx2char.get(next_char_idx, ""))
dec_input = np.array([next_char_idx], dtype=np.int64)
return "".join(output_chars)
def transliterate_sentence(self, sentence):
words = sentence.split()
translated_words = []
for word in words:
clean_word = "".join(c for c in word if '\u0900' <= c <= '\u097f')
if not clean_word:
translated_words.append(word)
continue
translated = self.transliterate_word(clean_word)
prefix = ""
for c in word:
if not ('\u0900' <= c <= '\u097f'):
prefix += c
else:
break
suffix = ""
for c in reversed(word):
if not ('\u0900' <= c <= '\u097f'):
suffix = c + suffix
else:
break
translated_words.append(prefix + translated + suffix)
return " ".join(translated_words)
def main():
print("=" * 80)
print("EVALUATING UPDATED TEXT-ALIGNED ONNX MODEL ON REAL TRANSCRIPTS")
print("=" * 80)
base_dir = os.path.dirname(os.path.abspath(__file__))
onnx_dir = os.path.join(base_dir, "../models")
transcript_file = os.path.join(base_dir, "./eval_transcript.json")
# 1. Load ONNX model
onnx_rnn = ONNXTransliterator(onnx_dir)
print("✓ ONNX RNN model loaded successfully!")
# 2. Load transcript segments
with open(transcript_file, "r", encoding="utf-8") as f:
data = json.load(f)
segments = data.get("segments", [])
test_sentences = []
for seg in segments:
text = seg.get("text", "").strip()
if any('\u0900' <= char <= '\u097f' for char in text):
test_sentences.append(text)
if len(test_sentences) >= 30:
break
# 3. Transliterate and print side-by-side
print("\n" + "=" * 90)
print(f"{'Devanagari Transcript Segment':<45} | {'Texting-Aligned Hinglish Output':<45}")
print("=" * 90)
latencies = []
for idx, sentence in enumerate(test_sentences):
t_start = time.perf_counter()
output = onnx_rnn.transliterate_sentence(sentence)
latency = (time.perf_counter() - t_start) * 1000 # ms
latencies.append(latency)
print(f"{idx+1:2d}. {sentence}")
print(f" -> '{output}' ({latency:.1f} ms)")
print("-" * 90)
print(f"\nAverage Sentence Transliteration Latency: {sum(latencies)/len(latencies):.2f} ms")
print("=" * 90)
if __name__ == "__main__":
main()
|