pilarhidalgo / find_span Goto Github PK

class SpanFinder:
    def __init__(self, io, sheet_name):
        self.io = io
        self.sheet_name = sheet_name
        self.patterns = self.read_patterns()
        self.nlp = self.set_nlp()
        self.trans = self.make_trans()
        
    def read_patterns(self):
        df = pd.read_excel(self.io, sheet_name=self.sheet_name)#puedes tener varios sheets con diferentes diccionarios en una sola hoja excel
        #creardo diccionarios en la estructura que Spacy requiere
        df1 = df[['label', 'pattern']]
        df2 = df[['label', 'pattern.1']]
        df2 = df2.rename(columns={'pattern.1': 'pattern'}, inplace=False)
        #Crear los diccionarios #si los nombres del conjunto SPAN son mas de 2 es necesario tambien implementar 
        #mas patterns
        patterns1 = df1.to_dict(orient='records')
        patterns2 = df2.to_dict(orient='records')
        #patternsn=dfn.to_dict(orient='records')
        patterns = patterns1 + patterns2
        return patterns
    
    def set_nlp(self):
        #Crear el objeto NLP
        nlp = Spanish()
        ruler = EntityRuler(nlp)
        ruler.add_patterns(self.patterns)
        nlp.add_pipe(ruler)
        return nlp
    
    def make_trans(self):
        #Preprocesar la entrada
        a,b = 'áéíóúü','aeiouu'
        trans = str.maketrans(a,b)
        return trans
    
    def find_span(self, sentence):
        sentence=sentence.lower().translate(self.trans)
        doc = self.nlp(sentence)
        #Devuelve la etiqueta real de la entidad
        if (sentence=='' or  sentence==' ' or sentence=='Por definir' or sentence=='por definir' or
           sentence=='no hay datos ' or sentence=='nan'):
            label=['No encontrado']#con esto evitaremos el rpoblema de comillas
        else: 
            label= list(set([(ent.text,ent.label_) for ent in doc.ents]))
        #label no establecida
        if label!=[]:
            label=label
        else: label=['No encontrado']#con esto evitaremos el problema de comillas
            
        return label[0]