import pickle
from utils.ImageAndTextEmbedding.index import getTextEmbedding

with open("word2vec_model.pkl", "rb") as f:
    textEmbedding_model = pickle.load(f)

def get_text_vector(example_text):
    # Tokenize the text into words
    words = example_text.lower().split()

    # Filter out words that are not in the vocabulary of the Word2Vec model
    words_in_vocab = [word for word in words if word in textEmbedding_model]

    # Calculate the average vector representation of the words
    if words_in_vocab:
        text_vector = sum(textEmbedding_model[word] for word in words_in_vocab) / len(words_in_vocab)
        return text_vector.tolist()
    else:
        return None

def get_text_discription_vector(text):
    return getTextEmbedding(text)

# Example usage:
# example_text = "This is an example sentence."
# text_vector = get_text_vector(example_text)
# if text_vector:
#     print("Vector representation of the example text:", text_vector)
# else:
#     print("None of the words in the example text are in the vocabulary of the Word2Vec model.")

print("Text embedding model loaded successfully!")