# Copyright (c) 2026 Edison dos Santos
# Licensed under the Apache License 2.0 (see LICENSE file in root)
# Part of Ghost Assistant: https://github.com/Edison2ST/GhostAssistantONNXFiles

import onnxruntime as ort
import numpy as np
import json
import librosa

with open("vocab.json", "r", encoding="utf-8") as f:
    vocab = json.load(f)

print("Loading ONNX model...")
session = ort.InferenceSession("omni_v2_q4.onnx")

print("Processing audio.wav...")
speech, _ = librosa.load("audio.wav", sr=16000, offset=0, duration=30)
speech = (speech - np.mean(speech)) / (np.std(speech) + 1e-7)
audio_input = np.expand_dims(speech, axis=0).astype(np.float32)

# Run ONNX Inference
logits = session.run(None, {session.get_inputs()[0].name: audio_input})[0]
predicted_ids = np.argmax(logits, axis=-1)[0]

# CTC Collapse & String Construction
transcript = ""
previous_id = -1

# Identify the SentencePiece space character (Lower One Eighth Block)
SP_SPACE = " " 

for idx in predicted_ids:
    idx = int(idx)
    if idx != previous_id:
        if idx != 0:  # Skip CTC blank/padding
            token = vocab.get(str(idx), "")
            
            # Ensure token is a string before calling .replace
            token_str = str(token)
            
            if token_str == SP_SPACE:
                transcript += " "
            else:
                # Clean up the SentencePiece marker and add to result
                transcript += token_str.replace(SP_SPACE, " ")
                
    previous_id = idx



print("\n" + "="*40)
print(f"TRANSCRIPTION: \n{transcript.strip().lower()}")
print("="*40)