forked from kubapok/retroc2
21 lines
715 B
Python
21 lines
715 B
Python
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.linear_model import LinearRegression
|
|
from stop_words import get_stop_words
|
|
import pandas as pd
|
|
|
|
|
|
|
|
vectorizer = TfidfVectorizer(stop_words=get_stop_words('polish'))
|
|
reg = LinearRegression()
|
|
colnames = ['start_date', 'end_date', 'title', 'sort_title', 'data']
|
|
t = pd.read_csv('test-A/in.tsv', sep="\t", names=colnames)
|
|
df = pd.read_csv('train/train.tsv', sep="\t", names=colnames)
|
|
date = (df['start_date'] + df['end_date']) / 2
|
|
train = vectorizer.fit_transform(df['data'])
|
|
reg.fit(train, date)
|
|
r = vectorizer.transform(t['data'])
|
|
pred = reg.predict(r)
|
|
with open('test-A/out.tsv', 'w') as file:
|
|
for e in pred:
|
|
file.write("%f\n" % e)
|