ireland-news-headlines/run.ipynb
2022-06-07 01:00:10 +02:00

123 lines
3.9 KiB
Plaintext

{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "26d23c74-6b2a-469d-9b26-3cc3f2ab32ba",
"metadata": {},
"outputs": [],
"source": [
"import vowpalwabbit\n",
"import pandas as pd\n",
"import re\n",
"\n",
"\n",
"x_train = pd.read_csv('train/in.tsv', header=None, sep='\\t')\n",
"y_train = pd.read_csv('train/expected.tsv', header=None, sep='\\t')\n",
"\n",
"x_train = x_train.drop(1, axis=1)\n",
"x_train.columns = ['year', 'text']\n",
"y_train.columns = ['category']\n",
"\n",
"data = pd.concat([x_train, y_train], axis=1)\n",
"\n",
"\n",
"model = vowpalwabbit.Workspace('--oaa 7 --ngram 3')\n",
"\n",
"map_dict = {}\n",
"\n",
"for i, x in enumerate(data['category'].unique()):\n",
" map_dict[x] = i+1 \n",
" \n",
"data['train_input'] = data.apply(lambda row: to_vw_format(row, map_dict), axis=1)\n",
"\n",
"\n",
"for example in data['train_input']:\n",
" model.learn(example)\n",
"\n",
"\n",
"\n",
"def to_vw_format(row, map_dict):\n",
" text = row['text'].replace('\\n', ' ').lower().strip()\n",
" text = re.sub(\"[^a-zA-Z -']\", '', text)\n",
" year = row['year']\n",
" try:\n",
" category = map_dict[row['category']]\n",
" except KeyError:\n",
" category = ''\n",
"\n",
" vw_input = f\"{category} | year:{year} text:{text}\\n\"\n",
"\n",
" return vw_input\n",
"\n",
"\n",
"\n",
"### Read data \n",
"\n",
"data_dev = pd.read_csv('dev-0/in.tsv', header=None, sep='\\t')\n",
"data_dev = data_dev.drop(1, axis=1)\n",
"data_dev.columns = ['year', 'text']\n",
"data_dev['train_input'] = data_dev.apply(lambda row: to_vw_format(row, map_dict), axis=1)\n",
"\n",
"\n",
"data_A = pd.read_csv('test-A/in.tsv', header=None, sep='\\t')\n",
"data_A = data_A.drop(1, axis=1)\n",
"data_A.columns = ['year', 'text']\n",
"data_A['train_input'] = data_A.apply(lambda row: to_vw_format(row, map_dict), axis=1)\n",
"\n",
"data_B = pd.read_csv('test-B/in.tsv', header=None, sep='\\t')\n",
"data_B = data_B.drop(1, axis=1)\n",
"data_B.columns = ['year', 'text']\n",
"data_B['train_input'] = data_B.apply(lambda row: to_vw_format(row, map_dict), axis=1)\n",
"\n",
"\n",
"\n",
"### Write predictions \n",
"\n",
"with open(\"dev-0/out.tsv\", 'w', encoding='utf-8') as file:\n",
" for test_example in data_dev['train_input']:\n",
" prediction_dev = model.predict(test_example)\n",
" text_prediction_dev = dict((value, key) for key, value in map_dict.items()).get(prediction_dev)\n",
" file.write(str(text_prediction_dev) + '\\n')\n",
"\n",
"\n",
"with open(\"test-A/out.tsv\", 'w', encoding='utf-8') as file:\n",
" for test_example in data_A['train_input']:\n",
" prediction_A = model.predict(test_example)\n",
" text_prediction_A = dict((value, key) for key, value in map_dict.items()).get(prediction_A)\n",
" file.write(str(text_prediction_A) + '\\n')\n",
"\n",
"\n",
"with open(\"test-B/out.tsv\", 'w', encoding='utf-8') as file:\n",
" for test_example in data_B['train_input']:\n",
" prediction_B = model.predict(test_example)\n",
" text_prediction_B = dict((value, key) for key, value in map_dict.items()).get(prediction_B)\n",
" file.write(str(text_prediction_B) + '\\n')\n",
"\n",
"\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.7"
}
},
"nbformat": 4,
"nbformat_minor": 5
}