models/text_encoder.py

import onnxruntime as ort
from utils.config import TEXT_ENCODER
from models.tokenizer import ModelTokenizer


class TextEncoder:
    def __init__(self):
        self.tokenizer = ModelTokenizer()
        self.text_model = ort.InferenceSession(TEXT_ENCODER)

    def get_text_embedding(self, text):
        inputs = self.tokenizer.tokenize_text(text)
        embedding = self.text_model.run(None, inputs)[0].squeeze(0)
        return embedding