utils.py

# -*- coding: utf-8 -*-
"""
@author: bai
"""
import logging, argparse


def get_entity(tag_seq, char_seq):
    PER = get_PER(tag_seq, char_seq)
    LOC = get_LOC(tag_seq, char_seq)
    ORG = get_ORG(tag_seq, char_seq)
    return PER, LOC, ORG


def get_PER(tag_seq, char_seq):
    length = len(char_seq)
    PER = []
    per=''
    for i, (char, tag) in enumerate(zip(char_seq, tag_seq)):
        if tag == 'B-PER':
            if per !='':
                PER.append(per)
                per=''
            per += char
            if i+1 == length:
                PER.append(per)
        if tag == 'I-PER':
            per += char
            if i+1 == length:
                PER.append(per)
        if tag not in ['I-PER', 'B-PER']:
            if per !='':
                PER.append(per)
                per=''
    return PER


def get_LOC(tag_seq, char_seq):
    length = len(char_seq)
    LOC = []
    loc=''
    for i, (char, tag) in enumerate(zip(char_seq, tag_seq)):
        if tag == 'B-LOC':
            if loc !='':
                LOC.append(loc)
                loc=''
            loc += char
            if i+1 == length:
                LOC.append(loc)
        if tag == 'I-LOC':
            loc += char
            if i+1 == length:
                LOC.append(loc)
        if tag not in ['I-LOC', 'B-LOC']:
            if loc !='':
                LOC.append(loc)
                loc=''
    return LOC


def get_ORG(tag_seq, char_seq):
    length = len(char_seq)
    ORG = []
    org=''
    for i, (char, tag) in enumerate(zip(char_seq, tag_seq)):
        if tag == 'B-ORG':
            if org !='':
               ORG.append(org)
               org=''
            org += char
            if i+1 == length:
                ORG.append(org)
        if tag == 'I-ORG':
            org += char
            if i+1 == length:
                ORG.append(org)
        if tag not in ['I-ORG', 'B-ORG']:
            if org !='':
                ORG.append(org)
                org=''
    return ORG

def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')
        
def get_logger(filename):
    logger = logging.getLogger('logger')
    logger.setLevel(logging.DEBUG)
    logging.basicConfig(format='%(message)s', level=logging.DEBUG)
    handler = logging.FileHandler(filename)
    handler.setLevel(logging.DEBUG)
    handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s: %(message)s'))
    logging.getLogger().addHandler(handler)
    return logger