from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LinearRegression
import sys
import pandas as pd

train_file = sys.argv[1]
pred_file = sys.argv[2]

train = pd.read_csv(train_file, sep='\t', header=None)
#pred_x = pd.read_csv(pred_file, sep='\t', header=None)
pred_x = []
with open(pred_file, encoding='utf-8') as f:
    for line in f:
        pred_x.append(line)
train_x, train_y = train[4], train[0]
#pred_x = pred[4]
#pred_x = pred_x.stack()

vectorizer = TfidfVectorizer()
train_x = vectorizer.fit_transform(train_x)
pred_x = vectorizer.transform(pred_x)

model = LinearRegression()
model.fit(train_x, train_y)

pred_y = model.predict(pred_x)

pd.DataFrame(pred_y).to_csv('out.tsv', header=False, index=None)