forked from kubapok/auta-public
43 lines
1.1 KiB
Python
43 lines
1.1 KiB
Python
import numpy
|
|
import pandas
|
|
import sys
|
|
from sklearn.linear_model import LinearRegression
|
|
|
|
TRAIN_URL = "train/train.tsv"
|
|
TEST_A = "test-A"
|
|
DEV_0 = "dev-0"
|
|
|
|
def get_trained_model(data):
|
|
x = numpy.array(data.get('x')).reshape(-1, 1)
|
|
y = numpy.array(data.get('y')).reshape(-1, 1)
|
|
m = LinearRegression()
|
|
m.fit(x, y)
|
|
return m
|
|
|
|
|
|
def prepare_model():
|
|
csv = pandas.read_csv(TRAIN_URL, sep='\t', header=None, index_col=None)
|
|
data = {'x': csv[2].tolist(), 'y': csv[0].tolist()}
|
|
model = get_trained_model(data)
|
|
return model
|
|
|
|
|
|
def calculate_prediction(path, prepared_model):
|
|
data_frame = pandas.read_csv(path + '/in.tsv', sep='\t', header=None, index_col=None)
|
|
input_array = numpy.array(data_frame[1].tolist()).reshape(-1, 1)
|
|
prediction = prepared_model.predict(input_array)
|
|
numpy.savetxt(path + "/out.tsv", prediction, fmt="%d", delimiter='\n')
|
|
|
|
|
|
def predict_dev_zero(prepared_model):
|
|
calculate_prediction(DEV_0, prepared_model)
|
|
|
|
|
|
def predict_test_a(prepared_model):
|
|
calculate_prediction(TEST_A, prepared_model)
|
|
|
|
|
|
prepared_model = prepare_model()
|
|
predict_dev_zero(prepared_model)
|
|
predict_test_a(prepared_model)
|