Ireland news headlines

run.py

@@ -0,0 +1,59 @@
+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)