Some changes

cooccurrence_count.py

@@ -21,6 +21,7 @@ import evaluate
 import datasets
 from sklearn.feature_extraction.text import CountVectorizer
 import numpy as np
+import stanza
 
 
 _DESCRIPTION = """\
@@ -53,6 +54,18 @@ def check_count(x):
     return 1
 
 
+nlp = stanza.Pipeline(lang='en', processors='tokenize')
+
+
+def stanza_tokenizer(sen):
+    doc = nlp(sen)
+    tokens = []
+    for sen in doc.sentences:
+        for token in sen.tokens:
+            tokens.append(token.text)
+    return tokens
+
+
 @evaluate.utils.file_utils.add_start_docstrings(
     _DESCRIPTION,
     _KWARGS_DESCRIPTION
@@ -71,9 +84,12 @@ class CooccurrenceCount(evaluate.Measurement):
             }),
         )
 
+    def _download_and_prepare(self, dl_manager):
+        stanza.download('en', processors='tokenize')
+
     def _compute(self, data, word1, word2):
-        len1 = len(word1.split(' '))
-        len2 = len(word2.split(' '))
+        len1 = len(stanza_tokenizer(word1))
+        len2 = len(stanza_tokenizer(word2))
         if len1 > len2:
             ugram = len1
             lgram = len2
@@ -84,14 +100,20 @@ class CooccurrenceCount(evaluate.Measurement):
             ugram = len1
             lgram = len1
 
-        v = CountVectorizer(ngram_range=(lgram, ugram))
+        v = CountVectorizer(
+            ngram_range=(lgram, ugram),
+            tokenizer=stanza_tokenizer,
+            lowercase=True
+        )
         analyzer = v.build_analyzer()
         vectorizer = CountVectorizer(
             ngram_range=(lgram, ugram),
             vocabulary={
                 analyzer(word1)[-1]: 0,
                 analyzer(word2)[-1]: 1
-            }
+            },
+            tokenizer=stanza_tokenizer,
+            lowercase=True
         )
         co_occurrences = vectorizer.fit_transform(data)
         dense_mat = co_occurrences.todense()