38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
from sklearn.linear_model import LinearRegression
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.metrics import mean_squared_error
|
|
import lzma
|
|
import numpy as np
|
|
|
|
X_train_raw = lzma.open("train/train.tsv.xz", mode='rt').readlines()
|
|
X_dev0 = open("dev-0/in.tsv", "r").readlines()
|
|
y_expected_dev0 = open("dev-0/expected.tsv", "r").readlines()
|
|
X_dev1 = open("dev-1/in.tsv", "r").readlines()
|
|
y_expected_dev1 = open("dev-1/expected.tsv", "r").readlines()
|
|
X_test = open("test-A/in.tsv", "r").readlines()
|
|
X = [i.split('\t') for i in X_train_raw]
|
|
X_train = [x[4] for x in X]
|
|
y_expected_train = [x[0] for x in X]
|
|
|
|
vectorizer = TfidfVectorizer(max_features=10000)
|
|
X_train_tfidf = vectorizer.fit_transform(X_train)
|
|
|
|
X_dev0_tfidf = vectorizer.transform(X_dev0)
|
|
X_dev1_tfidf = vectorizer.transform(X_dev1)
|
|
X_test_tfidf = vectorizer.transform(X_test)
|
|
|
|
model = LinearRegression()
|
|
model.fit(X_train_tfidf, y_expected_train)
|
|
y_predicted_dev0 = model.predict(X_dev0_tfidf)
|
|
y_predicted_dev1 = model.predict(X_dev1_tfidf)
|
|
y_predicted_test = model.predict(X_test_tfidf)
|
|
|
|
rmse_dev0 = mean_squared_error(y_expected_dev0, y_predicted_dev0, squared=False)
|
|
rmse_dev1 = mean_squared_error(y_expected_dev1, y_predicted_dev1, squared=False)
|
|
|
|
print(rmse_dev0, rmse_dev1)
|
|
|
|
open("dev-0/out.tsv", mode='w').writelines([str(i)+'\n' for i in y_predicted_dev0])
|
|
open("dev-1/out.tsv", mode='w').writelines([str(i)+'\n' for i in y_predicted_dev1])
|
|
open("test-A/out.tsv", mode='w').writelines([str(i)+'\n' for i in y_predicted_test])
|