forked from kubapok/retroc2
57 lines
1.8 KiB
Python
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()
|