89 lines
2.5 KiB
Python
89 lines
2.5 KiB
Python
import vowpalwabbit
|
|
import pandas as pd
|
|
import re
|
|
|
|
|
|
x_train = pd.read_csv('train/in.tsv', header=None, sep='\t')
|
|
y_train = pd.read_csv('train/expected.tsv', header=None, sep='\t')
|
|
|
|
x_train = x_train.drop(1, axis=1)
|
|
x_train.columns = ['year', 'text']
|
|
y_train.columns = ['category']
|
|
|
|
data = pd.concat([x_train, y_train], axis=1)
|
|
|
|
|
|
model = vowpalwabbit.Workspace('--oaa 7 --ngram 3')
|
|
|
|
map_dict = {}
|
|
|
|
for i, x in enumerate(data['category'].unique()):
|
|
map_dict[x] = i+1
|
|
|
|
data['train_input'] = data.apply(lambda row: to_vw_format(row, map_dict), axis=1)
|
|
|
|
|
|
for example in data['train_input']:
|
|
model.learn(example)
|
|
|
|
|
|
|
|
def to_vw_format(row, map_dict):
|
|
text = row['text'].replace('\n', ' ').lower().strip()
|
|
text = re.sub("[^a-zA-Z -']", '', text)
|
|
year = row['year']
|
|
try:
|
|
category = map_dict[row['category']]
|
|
except KeyError:
|
|
category = ''
|
|
|
|
vw_input = f"{category} | year:{year} text:{text}\n"
|
|
|
|
return vw_input
|
|
|
|
|
|
|
|
### Read data
|
|
|
|
data_dev = pd.read_csv('dev-0/in.tsv', header=None, sep='\t')
|
|
data_dev = data_dev.drop(1, axis=1)
|
|
data_dev.columns = ['year', 'text']
|
|
data_dev['train_input'] = data_dev.apply(lambda row: to_vw_format(row, map_dict), axis=1)
|
|
|
|
|
|
data_A = pd.read_csv('test-A/in.tsv', header=None, sep='\t')
|
|
data_A = data_A.drop(1, axis=1)
|
|
data_A.columns = ['year', 'text']
|
|
data_A['train_input'] = data_A.apply(lambda row: to_vw_format(row, map_dict), axis=1)
|
|
|
|
data_B = pd.read_csv('test-B/in.tsv', header=None, sep='\t')
|
|
data_B = data_B.drop(1, axis=1)
|
|
data_B.columns = ['year', 'text']
|
|
data_B['train_input'] = data_B.apply(lambda row: to_vw_format(row, map_dict), axis=1)
|
|
|
|
|
|
|
|
### Write predictions
|
|
|
|
with open("dev-0/out.tsv", 'w', encoding='utf-8') as file:
|
|
for test_example in data_dev['train_input']:
|
|
prediction_dev = model.predict(test_example)
|
|
text_prediction_dev = dict((value, key) for key, value in map_dict.items()).get(prediction_dev)
|
|
file.write(str(text_prediction_dev) + '\n')
|
|
|
|
|
|
with open("test-A/out.tsv", 'w', encoding='utf-8') as file:
|
|
for test_example in data_A['train_input']:
|
|
prediction_A = model.predict(test_example)
|
|
text_prediction_A = dict((value, key) for key, value in map_dict.items()).get(prediction_A)
|
|
file.write(str(text_prediction_A) + '\n')
|
|
|
|
|
|
with open("test-B/out.tsv", 'w', encoding='utf-8') as file:
|
|
for test_example in data_B['train_input']:
|
|
prediction_B = model.predict(test_example)
|
|
text_prediction_B = dict((value, key) for key, value in map_dict.items()).get(prediction_B)
|
|
file.write(str(text_prediction_B) + '\n')
|
|
|
|
|