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

# cd ..

import tagnews
import pandas as pd

# Download (and extract if needed) a saved glove data from
# https://github.com/stanfordnlp/GloVe
# and save it to tagnews/data/
glove = tagnews.load_glove('tagnews/data/glove.6B.50d.txt')

# Download (and extract if needed) the NER data from
# https://www.kaggle.com/abhinavwalia95/entity-annotated-corpus/data
# and save it to tagnews/data/
ner = tagnews.load_ner_data('tagnews/data/')

ner = pd.concat([ner, pd.DataFrame(glove.loc[ner['word'].str.lower()].values)], axis='columns')

num_asserts = 0
for i, row in ner.sample(1000).iterrows():
    if not any(row.iloc[2:].isnull()):
        assert (glove.loc[row['word'].lower()].values == row.iloc[3:].values).all()
        num_asserts += 1
print('Asserted correct vectorizations', num_asserts, 'times.')

import sklearn.ensemble

clf = sklearn.ensemble.RandomForestClassifier()

# be careful doing this if you are relying on sequential-ness!
ner.fillna(value=0.0, inplace=True)

clf.fit(ner.iloc[:200000, 3:], ner['tag'].iloc[:200000].values)

clf.predict_proba(glove.loc[['london', 'france', 'napkins']])

# Go to https://geo-extract-tester.herokuapp.com/ and download
# the validation data (validation.txt).
with open('validation.txt', encoding='utf-8') as f:
    s = f.read()

with open('guesses.txt', 'w') as f:
    for prob in clf.predict_proba(glove.loc[[w for w in s.split('\n') if w]].fillna(0))[:,1]:
        f.write(str(prob) + '\n')

# Now go to https://geo-extract-tester.herokuapp.com/ and upload `guesses.txt` to see how you did!



StarcoderdataJupyterScriptsDedupFiltered