# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# # !wget https://www.dropbox.com/s/lnbhc8yuhit4wm5/wikidata5m_alias.tar.gz?dl=1
# # !mv wikidata5m_alias.tar.gz?dl=1 wikidata5m_alias.tar.gz
# # !tar -zxf wikidata5m_alias.tar.gz
# # !rm wikidata5m_alias.tar.gz
# -

# !head -n 100 wikidata5m_entity.txt

# +
with open('wikidata5m_entity.txt') as fopen:
    data = fopen.read().split('\n')
    
data[0]

# +
import re
from unidecode import unidecode
from collections import defaultdict
import numpy as np

def cleaning(string):
    string = re.sub('[^A-Za-z ]+', ' ', unidecode(string).lower())
    return re.sub(r'[ ]+', ' ', string).strip()


# -

def get_uniques(string):
    uniques = defaultdict(list)
    splitted = string.split('\t')[1:]
    for s in splitted:
        uniques[cleaning(s)].append(unidecode(s))
        
    r = []  
    for k, v in uniques.items():
        vals = [sum([ord(c) for c in v_]) for v_ in v]
        r.append(v[np.argmin(vals)])
    return r


# %%time
get_uniques(data[0])

len(data)

# +
import ujson as json

filename = 'wikidata5m_entity.txt'
partition_size = 10
limit = len(data) // partition_size
for i in range(partition_size):
    b = data[i * limit: limit + (i * limit)]
    b = [r for r in b if len(r)]
    with open(f'{filename}-{i}', 'w') as fopen:
        json.dump(b, fopen)
# -



StarcoderdataJupyterScriptsDedupFiltered