95 lines
2.0 KiB
Python
95 lines
2.0 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
# In[3]:
|
|
|
|
|
|
import vowpalwabbit
|
|
import pandas as pd
|
|
import re
|
|
|
|
|
|
# In[4]:
|
|
|
|
|
|
def prediction(path_in, path_out, model, categories):
|
|
data = pd.read_csv(path_in, header=None, sep='\t')
|
|
data = data.drop(1, axis=1)
|
|
data.columns = ['year', 'text']
|
|
|
|
data['train_input'] = data.apply(lambda row: to_vowpalwabbit(row, categories), axis=1)
|
|
|
|
with open(path_out, 'w', encoding='utf-8') as file:
|
|
for example in data['train_input']:
|
|
predicted = model.predict(example)
|
|
text_predicted = dict((value, key) for key, value in categories.items()).get(predicted)
|
|
file.write(str(text_predicted) + '\n')
|
|
|
|
|
|
# In[5]:
|
|
|
|
|
|
def to_vowpalwabbit(row, categories):
|
|
text = row['text'].replace('\n', ' ').lower().strip()
|
|
text = re.sub("[^a-zA-Z -']", '', text)
|
|
text = re.sub(" +", ' ', text)
|
|
year = row['year']
|
|
try:
|
|
category = categories[row['category']]
|
|
except KeyError:
|
|
category = ''
|
|
|
|
vw = f"{category} | year:{year} text:{text}\n"
|
|
|
|
return vw
|
|
|
|
|
|
# In[10]:
|
|
|
|
|
|
x_train = pd.read_csv('train/in.tsv', header=None, sep='\t')
|
|
x_train = x_train.drop(1, axis=1)
|
|
x_train.columns = ['year', 'text']
|
|
|
|
|
|
y_train = pd.read_csv('train/expected.tsv', header=None, sep='\t')
|
|
y_train.columns = ['category']
|
|
|
|
x_train = x_train[0:800000]
|
|
y_train = y_train(0:800000)
|
|
|
|
data = pd.concat([x_train, y_train], axis=1)
|
|
|
|
|
|
# In[ ]:
|
|
|
|
|
|
categories = {}
|
|
|
|
for i, x in enumerate(data['category'].unique()):
|
|
categories[x] = i+1
|
|
|
|
print(categories)
|
|
|
|
data['train_input'] = data.apply(lambda row: to_vowpalwabbit(row, categories), axis=1)
|
|
|
|
model = vowpalwabbit.Workspace('--oaa 7 --learning_rate 0.99')
|
|
|
|
for example in data['train_input']:
|
|
model.learn(example)
|
|
|
|
|
|
# In[ ]:
|
|
|
|
|
|
prediction('dev-0/in.tsv', 'dev-0/out.tsv', model, categories)
|
|
prediction('test-A/in.tsv', 'test-A/out.tsv', model, categories)
|
|
prediction('test-B/in.tsv', 'test-B/out.tsv', model, categories)
|
|
|
|
|
|
# In[ ]:
|
|
|
|
|
|
get_ipython().system('jupyter nbconvert --to script run.ipynb')
|
|
|