import sys
import json
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"


def create_embedding(text: str) -> list[float]:
    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)

    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)

    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]

    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

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

    return embedding[0].tolist()


if __name__ == "__main__":
    text = sys.stdin.read().strip()

    if not text:
        print(json.dumps({
            "error": "No text supplied"
        }))
        sys.exit(1)

    embedding = create_embedding(text)

    print(json.dumps({
        "dimensions": len(embedding),
        "embedding": embedding,
    }))