<commit_before>from nalaf.features import FeatureGenerator
from gensim.models import Word2Vec


class WordEmbeddingsFeatureGenerator(FeatureGenerator):
    """
    DOCSTRING
    """

    def __init__(self, model_file, weight=1):
        self.model = Word2Vec.load(model_file)
        self.weight = weight

    def generate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        for token in dataset.tokens():
            if token.word.lower() in self.model:
                for index, value in enumerate(self.model[token.word.lower()]):
                    # value.item() since value is a numpy float
                    # and we want native python floats
                    token.features['embedding_{}'.format(index)] = self.weight * value.item()
<commit_msg>features based on brown clusters<commit_after>from nalaf.features import FeatureGenerator
from gensim.models import Word2Vec


class WordEmbeddingsFeatureGenerator(FeatureGenerator):
    """
    DOCSTRING
    """

    def __init__(self, model_file, weight=1):
        self.model = Word2Vec.load(model_file)
        self.weight = weight

    def generate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        for token in dataset.tokens():
            if token.word.lower() in self.model:
                for index, value in enumerate(self.model[token.word.lower()]):
                    # value.item() since value is a numpy float
                    # and we want native python floats
                    token.features['embedding_{}'.format(index)] = self.weight * value.item()


class BrownClusteringFeatureGenerator(FeatureGenerator):
        """
        DOCSTRING
        """
        def __init__(self, model_file, weight=1):
            with open(model_file, encoding='utf-8') as file:
                self.clusters = {str(line.split()[1]): line.split()[0] for line in file.readlines()}
            self.weight = weight

        def generate(self, dataset):
            """
            :type dataset: nalaf.structures.data.Dataset
            """
            for token in dataset.tokens():
                if token.word in self.clusters:
                    assignment = self.clusters[token.word]
                    for i in range(len(assignment)):
                        token.features['brown'] = assignment[:i+1]

{'fraction_non_alphanumeric': 0.10723296888141295, 'fraction_numerical': 0.004625735912531539, 'mean_word_length': 2.0896103896103897, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '333264', 'n_tokens_mistral': 643, 'n_tokens_neox': 611, 'n_words': 155}