utils_model.py

import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer

class ModelFactory():

    def __init__(self):
        pass

    def create_model(self, model_type):
        model = None

        if (model_type=='all-MiniLM-L6-v2'):
            model = MiniLM_L6_v2_Model()

        if (model_type=='sentence_similarity_spanish'):
            model = SentenceSimilaritySpanishModel()

        return model

class BaseModel():

    def __init__(self):
        pass

    def retrieve_embeddings(self, input_text):
        embeddings = self.model.encode(input_text, batch_size=32)
        embeddings *= 255
        embeddings = embeddings.astype(np.uint8).tolist()

        return embeddings

class MiniLM_L6_v2_Model(BaseModel):
    
    def __init__(self):
        self.model = SentenceTransformer('all-MiniLM-L6-v2')

class SentenceSimilaritySpanishModel(BaseModel):

    def __init__(self):
        self.model = SentenceTransformer('hiiamsid/sentence_similarity_spanish_es')