import gradio as gr
import numpy as np
import onnxruntime as ort
import sentencepiece as spm
import json
import os
import warnings
import asyncio

# 1. Подавление предупреждений asyncio
warnings.filterwarnings("ignore", category=RuntimeWarning, message="coroutine.*was never awaited")

# 2. Проверка наличия всех файлов
required_files = [
    "sentencepiece.bpe.model",
    "special_tokens_map.json",
    "encoder_model_quantized.onnx",
    "decoder_model_quantized.onnx"
]

for f in required_files:
    if not os.path.exists(f):
        raise FileNotFoundError(f"Критическая ошибка: файл не найден — {f}. Проверьте загрузку файлов в проект.")

# 3. Загрузка токенизатора
try:
    sp = spm.SentencePieceProcessor()
    sp.Load("sentencepiece.bpe.model")
except Exception as e:
    raise RuntimeError(f"Ошибка загрузки токенизатора: {e}")

# 4. Загрузка специальных токенов
try:
    with open("special_tokens_map.json", "r") as f:
        special_tokens = json.load(f)
except Exception as e:
    raise RuntimeError(f"Ошибка чтения special_tokens_map.json: {e}")

# 5. Загрузка ONNX-моделей
try:
    encoder_session = ort.InferenceSession("encoder_model_quantized.onnx")
    decoder_session = ort.InferenceSession("decoder_model_quantized.onnx")
except Exception as e:
    raise RuntimeError(f"Ошибка загрузки ONNX-модели: {e}")  # ЗАКРЫТА кавычка и скобка!


# 6. Отладка: вывод информации о моделях
print("\n=== ВХОДЫ ENCODER ===")
for inp in encoder_session.get_inputs():
    print(f"Имя: {inp.name}, Форма: {inp.shape}, Тип: {inp.type}")

print("\n=== ВЫХОДЫ ENCODER ===")
for out in encoder_session.get_outputs():
    print(f"Имя: {out.name}, Форма: {out.shape}, Тип: {out.type}")

print("\n=== ВХОДЫ DECODER ===")
for inp in decoder_session.get_inputs():
    print(f"Имя: {inp.name}, Форма: {inp.shape}, Тип: {inp.type}")