extraction.py

try:
    from .src.run import NLP_classification
except:
    from src.run import NLP_classification
import yaml
import os


class paragraph_extract:
    def __init__(self):
        config_file = 'configs/config.yaml'
        config_file = os.path.join(os.path.dirname(__file__), config_file)
        self.config = yaml.load(open(config_file), Loader=yaml.FullLoader)
        self.config['extract']['selected_model'] = os.path.join(os.path.dirname(__file__), self.config['extract']['selected_model'])
        
        self.runner = NLP_classification(**self.config['model'])

    def extract(self, paragraphs, positions):

        labels = self.runner.label_extraction(paragraphs, positions, **self.config['extract'])
        
        return labels