retroc2/run.py

29 lines
784 B
Python
Raw Permalink Normal View History

2022-04-28 22:47:06 +02:00
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)
2022-04-29 18:46:18 +02:00
#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)
2022-04-28 22:47:06 +02:00
train_x, train_y = train[4], train[0]
#pred_x = pred[4]
2022-04-29 18:46:18 +02:00
#pred_x = pred_x.stack()
2022-04-28 22:47:06 +02:00
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)