retroc2/run.py
2022-05-09 21:42:01 +02:00

57 lines
1.8 KiB
Python

#!/usr/bin/env python
import os
import math
import pandas as pd
from sklearn.metrics import mean_squared_error
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import make_pipeline
from stop_words import get_stop_words
def rmse(y_expected, y_predicted):
return math.sqrt(mean_squared_error(y_expected, y_predicted))
def read_tsv(path, filename, col_names, cutoff=None):
i = 0
data = []
with open(os.path.join(path, filename), 'r', encoding='utf-8') as file:
for line in file:
if cutoff and i >= cutoff:
break
data.append(line.strip().split('\t'))
i += 1
df = pd.DataFrame(data, columns=col_names)
return df
def evaluate_and_save(path, file, model):
df = read_tsv(path, file, ['y'])['y'].values
y_predicted = model.predict(df)
# y_expected = read_tsv(path, 'expected.tsv', ['y'])['y'].values
# print('score: ', rmse(y_expected, y_predicted))
with open(os.path.join(path, 'out.tsv'), 'w') as f:
for value in y_predicted:
f.write(f'{value}\n')
def main():
train_data = read_tsv('train', 'train.tsv', ['date_start', 'date_end', 'x1', 'x2', 'text'], 50000)
train_data['date'] = (train_data['date_start'].astype(float) + train_data['date_end'].astype(float)) / 2
train_x = train_data['text'].values
train_y = train_data['date'].values
vectorizer = TfidfVectorizer(stop_words=get_stop_words('polish'))
model = LinearRegression()
pipe = make_pipeline(vectorizer, model)
pipe.fit(train_x, train_y)
# print('training done')
for path, file in [('dev-0', 'in.tsv'), ('dev-1', 'in.tsv'), ('test-A', 'in.tsv')]:
evaluate_and_save(path, file, pipe)
if __name__ == '__main__':
main()