29 lines
784 B
Python
29 lines
784 B
Python
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)
|