modifies wikidata feature list

dibyaaaaax · dibyaaaaax · commit a07e48c31a30 · 2020-07-22T19:08:55.000+05:30
diff --git a/drafttopic/feature_lists/wikidata.py b/drafttopic/feature_lists/wikidata.py
@@ -1,23 +1,40 @@
+import re
+
+from revscoring.datasources import Datasource
 from revscoring.datasources.meta import vectorizers
 from revscoring.features import wikibase
 from revscoring.features.meta import aggregators
 
 
+QID_RE = re.compile('Q[0-9]+')
+
 wikidata_kvs = vectorizers.word2vec.load_gensim_kv(
     filename="wikidata-20200501-learned_vectors.50_cell.10k.kv", mmap="r")
 
+def process_claims_to_words(claims):
+    words = []
+    for pid, value in claims:
+        words.append(pid)
+        if QID_RE.match(value) is not None:
+            words.append(value)
+    return words
 
 def vectorize_words(words):
     return vectorizers.word2vec.vectorize_words(wikidata_kvs, words)
 
 
-revision_text_vectors = vectorizers.word2vec(
-    wikibase.revision.datasources.claims,
+claim_words = Datasource(
+    "wikidata.revision.claim_words",
+    process_claims_to_words,
+    depends_on=[wikibase.revision.datasources.claims])
+
+revision_claim_words_vectors = vectorizers.word2vec(
+    claim_words,
     vectorize_words,
     name="revision.text.wikidata_vectors")
 
 w2v = aggregators.mean(
-    revision_text_vectors,
+    revision_claim_words_vectors,
     vector=True,
     name="revision.text.wikidata_vectors_mean"
 )
diff --git a/drafttopic/utilities/balance_sample.py b/drafttopic/utilities/balance_sample.py
@@ -70,15 +70,17 @@ def balance_sample(observations, lang_code, max_n, output):
 def group_labeled_obs(observations, lang_code):
     label_obs = defaultdict(list)
     for ob in observations:
-        if lang_code in ob['sitelinks']:
-            labeled_ob = {
-                'title': ob['sitelinks'][lang_code]}
-        elif lang_code == 'wikidata':
-            labeled_ob = {
-                'title': ob['qid']}
-
-        if lang_code == 'wikidata' or lang_code in ob['sitelinks']:
-            labeled_ob['taxo_labels'] = ob['taxo_labels']
-            for taxo_label in ob['taxo_labels']:
-                label_obs[taxo_label].append(labeled_ob)
+        if lang_code == 'wikidata':
+            if ob['qid'] is None:
+                continue
+            labeled_ob = {'title': ob['qid']}
+        else:
+            if lang_code not in ob['sitelinks']:
+                continue
+            labeled_ob = {'title': ob['sitelinks'][lang_code]}
+
+        labeled_ob['taxo_labels'] = ob['taxo_labels']
+        for taxo_label in ob['taxo_labels']:
+            label_obs[taxo_label].append(labeled_ob)
+
     return label_obs
