prb977/cooccurrence_count

1# Copyright 2020 The HuggingFace Datasets Authors and the current2# dataset script contributor.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""16Get the co-occurance count for two words in each sentece in a dataset.17"""18 19 20import evaluate21import datasets22from sklearn.feature_extraction.text import CountVectorizer23import numpy as np24import stanza25 26 27_DESCRIPTION = """\28Returns the co-occurrence count of two words in the input.29"""30 31_CITATION = ""32 33_KWARGS_DESCRIPTION = """34Calculates the co-occurence of two words in each sentence.35Args:36    `data`: a list of `str` which containes a dataset.37    `words`: list of list of two words that we want to check for38Returns:39Examples:40    >>> data = ["hello sun","hello moon", "hello sun"]41    >>> c_count = evaluate.load("prb977/cooccurrence_count")42    >>> results = c_count.compute(data=data, words=[['hello','sun']\)43    >>> print(results)44    [['hello','sun',3,2]]45"""46 47 48def check_count(x):49    if x[0].all() <= 0:50        return 051    return 152 53 54nlp = stanza.Pipeline(lang='en', processors='tokenize')55 56 57def stanza_tokenizer(sen):58    doc = nlp(sen)59    tokens = []60    for sen in doc.sentences:61        for token in sen.tokens:62            tokens.append(token.text)63    return tokens64 65 66@evaluate.utils.file_utils.add_start_docstrings(67    _DESCRIPTION,68    _KWARGS_DESCRIPTION69)70class CooccurrenceCount(evaluate.Measurement):71    """This measurement returns the co-occurrence count of two words."""72 73    def _info(self):74        return evaluate.MeasurementInfo(75            module_type="measurement",76            description=_DESCRIPTION,77            citation=_CITATION,78            inputs_description=_KWARGS_DESCRIPTION,79            features=datasets.Features({80                'data': datasets.Value('string')81            }),82        )83 84    def _download_and_prepare(self, dl_manager):85        stanza.download('en', processors='tokenize')86 87    def _compute(self, data, words):88        for each in words:89            word1 = each[0]90            word2 = each[1]91            print(word1)92            print(word2)93            len1 = len(stanza_tokenizer(word1))94            len2 = len(stanza_tokenizer(word2))95            if len1 > len2:96                ugram = len197                lgram = len298            elif len1 < len2:99                ugram = len2100                lgram = len1101            else:102                ugram = len1103                lgram = len1104 105            v = CountVectorizer(106                ngram_range=(lgram, ugram),107                tokenizer=stanza_tokenizer,108                lowercase=True109            )110            analyzer = v.build_analyzer()111            vectorizer = CountVectorizer(112                ngram_range=(lgram, ugram),113                vocabulary={114                    analyzer(word1)[-1]: 0,115                    analyzer(word2)[-1]: 1116                },117                tokenizer=stanza_tokenizer,118                lowercase=True119            )120            co_occurrences = vectorizer.fit_transform(data)121            dense_mat = co_occurrences.todense()122            count = len(data)123            co_occurrence_count = np.sum(124                np.apply_along_axis(check_count, axis=1, arr=dense_mat)125            )126            each.append(count)127            each.append(co_occurrence_count)128        return words129