File size: 1,005 Bytes
9fa6c15
 
04178ea
9fa6c15
 
 
 
 
 
 
 
 
04178ea
 
9fa6c15
28032f0
 
 
9fa6c15
 
 
 
 
 
 
 
28032f0
 
 
9fa6c15
28032f0
9fa6c15
04178ea
9fa6c15
04178ea
 
9fa6c15
28032f0
9fa6c15
28032f0
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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')