agent/target_extraction/product.py

from anytree import Node
import pickle
from os.path import isfile
from agent.target_extraction.argument import Argument
from nltk.stem import WordNetLemmatizer

wnl = WordNetLemmatizer()


class Product:
    FILE_DIR = '/content/drive/MyDrive/EXAMPLE_DATASET/'
    FILE_EXTENSION = '.pickle'

    def __init__(self, root: Node, syn_dict):
        self.root = root
        self.feature_nodes = [n for n in root.descendants]
        self.argument_nodes = [root] + self.feature_nodes
        self.glossary = {a_node: [syn.split(' ') for syn in syns]
                         for a, syns in syn_dict.items() for a_node in self.argument_nodes if a_node.name == a}
        self.arguments = {a_node: Argument(a_idx, a_node.name.replace('_', ' '))
                          for a_idx, a_node in enumerate(self.argument_nodes)}
        self.singularities = {a_node: wnl.lemmatize(a_node.name) == a_node.name for a_node in self.argument_nodes}

    def argument_node_for_id(self, id):
        return self.argument_nodes[id]

    def argument_for_id(self, id):
        return self.argument_for_node(self.argument_node_for_id(id))

    def argument_for_node(self, n):
        return self.arguments[n]

    def name(self):
        return self.root.name

    @staticmethod
    def get_product(name):
        path = Product.FILE_DIR + name + Product.FILE_EXTENSION
        if isfile(path):
            f = open(path, 'rb')
            product: Product = pickle.load(f)
            f.close()
            return product
        else:
            raise Exception('No representation found for product {} at {}'.format(name, path))