import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

MODEL_PATH = "/var/www/html/manbot/embedding-model/onnx/model_qint8_arm64.onnx"
TOKENIZER_PATH = "/var/www/html/manbot/embedding-model/tokenizer.json"

text = """
What is the correct adjustment procedure for the take-up unit on the FC9000?
"""

print("Loading tokenizer...")
tokenizer = Tokenizer.from_file(TOKENIZER_PATH)

print("Tokenizing...")
encoded = tokenizer.encode(text)

input_ids = np.array([encoded.ids], dtype=np.int64)
attention_mask = np.array([encoded.attention_mask], dtype=np.int64)
token_type_ids = np.array([encoded.type_ids], dtype=np.int64)

print("Tokens:", len(encoded.ids))

print("Loading ONNX model...")
session = ort.InferenceSession(
    MODEL_PATH,
    providers=["CPUExecutionProvider"],
)

outputs = session.run(
    None,
    {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "token_type_ids": token_type_ids,
    },
)

token_embeddings = outputs[0]

# Attention-mask mean pooling
mask = attention_mask[..., None].astype(np.float32)

sum_embeddings = np.sum(token_embeddings * mask, axis=1)
sum_mask = np.clip(mask.sum(axis=1), a_min=1e-9, a_max=None)

embedding = sum_embeddings / sum_mask

# L2 normalization
embedding = embedding / np.linalg.norm(
    embedding,
    axis=1,
    keepdims=True
)

embedding = embedding[0]

print()
print("Embedding created")
print("Dimensions:", len(embedding))
print("Dtype:", embedding.dtype)
print("Norm:", np.linalg.norm(embedding))
print()
print("First 20 values:")
print(embedding[:20])
print()
print("TEST COMPLETED")