# ---
# 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
# ---

# %load_ext autoreload
# %autoreload 2

# %matplotlib inline

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
import seaborn as sns
sns.set_context('paper')
sns.set_style('white')

import sys
sys.path.append('/home/dan/bclm')

spdf = pd.read_csv('../data/spdf_fixed.csv.gz', low_memory=False)
spdf.head()

spdf.head().T

# +
from collections import OrderedDict
DEFAULT_FIELDS = ['id', 'form', 'lemma', 'upostag', 'xpostag', 'feats', 'head', 'deprel', 'deps', 'misc']
def sentence_to_conllu_token_list(sent, add_misc=['ner_escaped', 'token_id', 'token_str', 'biose_layer0', 
                                                  'biose_layer1', 'biose_layer2', 'biose_layer3'],
                                  metadata=['sent_id', 'global_sent_id', 'set', 'duplicate_sent_id', 'very_similar_sent_id']):
    metadata = OrderedDict([(m, str(sent[m].iat[0])) for m in metadata])
    
    if add_misc is not None:
        amisc = sent.apply(lambda x: [m+'='+str(x[m]) for m in add_misc], axis=1)
        amisc = amisc.str.join('|')
    token_morphs = (sent.groupby('token_id').id.agg([min, max])
                 .pipe(lambda x: x[x['max'] - x['min']>0])
                 .assign(token_morphs=lambda x: x['min'].astype(str)+'-'+x['max'].astype(str))
                 .set_index('min')['token_morphs'])   
    token_list = []
    
    for (i, row), token_str, am in zip(sent[DEFAULT_FIELDS].iterrows(), sent['token_str'].tolist(), amisc.tolist()):
        token = OrderedDict(row)
        if type(token['feats']) == str and token['feats']!='_':
            #token['feats'] = OrderedDict([f.split('=') for f in token['feats'].split('|')])
            feats = OrderedDict()
            for f in token['feats'].split('|'):
                k,v = f.split('=')
                if k not in feats:
                    feats[k] = v
                else:
                    feats[k] = feats[k]+','+v
            token['feats'] = feats
                        
        am = OrderedDict([f.split('=') for f in am.split('|') ])
        if token['misc'] =='_' or token['misc'] is None:
            token['misc'] = am
        elif type(token['misc'])==str:
            token['misc'] = OrderedDict([('MISC', token['misc'])])
            token['misc'].update(am)
        elif type(token['misc'])==OrderedDict:
            am.update(token['misc'])
            token['misc'] = am
            
        if token['id'] in token_morphs:
            tok_row = OrderedDict([(f, '_') for f in DEFAULT_FIELDS])
            id_tup = token_morphs[token['id']].split('-')
            tok_row['id'] = (id_tup[0], '-', id_tup[1])
            tok_row['form'] = token_str
            token_list.append(tok_row)
            
        token_list.append(token)
        
    return token_list, metadata


# -

from conllu.models import TokenList
sp_conllu_sents = [TokenList(*sentence_to_conllu_token_list(spdf[spdf.sent_id==i])) for i in spdf.sent_id.drop_duplicates().tolist()]
print(sp_conllu_sents[2].serialize())

# + jupyter={"outputs_hidden": true}
for s in sp_conllu_sents:
    for t in s:
        if type(t['feats'])==OrderedDict:
            for k, v in t['feats'].items():
                if ',' in v:
                    print (k, v)
# -



StarcoderdataJupyterScriptsDedupFiltered