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=='mock'):
            model = MockModel()

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

        return model

class BaseModel():

    def __init__(self):
        pass

    def retrieve_embeddings(self, input_text):
        pass    

class MockModel(BaseModel):

    def __init__(self):
        pass

    def retrieve_embeddings(self, input_text):
        random_embeddings = np.random.randint(256, size=(370))/256

        return pd.DataFrame(random_embeddings)

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

#     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