60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
|
import vowpalwabbit
|
||
|
import pandas as pd
|
||
|
import re
|
||
|
|
||
|
|
||
|
def to_vw_format(row, map_dict):
|
||
|
text = row['text'].replace('\n', ' ').lower().strip()
|
||
|
#text = re.sub("[^a-zA-Z0-9 -']", '', text)
|
||
|
text = re.sub("[^a-zA-Z -']", '', text)
|
||
|
text = re.sub(" +", ' ', 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
|
||
|
|
||
|
|
||
|
def predict_and_write(folder_name, model, map_dict):
|
||
|
data = pd.read_csv(f'{folder_name}/in.tsv', header=None, sep='\t')
|
||
|
data = data.drop(1, axis=1)
|
||
|
data.columns = ['year', 'text']
|
||
|
|
||
|
data['train_input'] = data.apply(lambda row: to_vw_format(row, map_dict), axis=1)
|
||
|
|
||
|
with open(f"{folder_name}/out.tsv", 'w', encoding='utf-8') as file:
|
||
|
for test_example in data['train_input']:
|
||
|
prediction = model.predict(test_example)
|
||
|
text_prediction = dict((value, key) for key, value in map_dict.items()).get(prediction)
|
||
|
file.write(str(text_prediction) + '\n')
|
||
|
|
||
|
|
||
|
model = vowpalwabbit.Workspace('--oaa 7')
|
||
|
|
||
|
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)
|
||
|
map_dict = {}
|
||
|
for i, x in enumerate(data['category'].unique()):
|
||
|
map_dict[x] = i+1 #0 nie może być
|
||
|
print(map_dict)
|
||
|
data['train_input'] = data.apply(lambda row: to_vw_format(row, map_dict), axis=1)
|
||
|
print(data.head(5))
|
||
|
|
||
|
|
||
|
for example in data['train_input']:
|
||
|
model.learn(example)
|
||
|
|
||
|
predict_and_write('dev-0', model, map_dict)
|
||
|
predict_and_write('test-A', model, map_dict)
|
||
|
predict_and_write('test-B', model, map_dict)
|