forked from kubapok/retroc2
30 lines
993 B
Python
30 lines
993 B
Python
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
||
|
|
from sklearn.linear_model import LinearRegression
|
||
|
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
import csv
|
||
|
|
|
||
|
|
def getData(path):
|
||
|
|
with open(path, encoding="utf-8") as source:
|
||
|
|
return source.readlines()
|
||
|
|
|
||
|
|
def changeToDf(input):
|
||
|
|
return pd.read_csv(input, sep="\t")
|
||
|
|
|
||
|
|
vectorizer = TfidfVectorizer()
|
||
|
|
linear = LinearRegression()
|
||
|
|
train = pd.read_csv("./train/train.tsv", sep="\t", names=['start_date', 'end_date', 'title', 'sort_title', 'data'])
|
||
|
|
mean = (train['start_date'] + train['end_date']) / 2
|
||
|
|
tv = vectorizer.fit_transform(train['data'])
|
||
|
|
linear.fit(tv, mean)
|
||
|
|
|
||
|
|
def calculateResult(in_, out):
|
||
|
|
tmp = getData(in_)
|
||
|
|
df = pd.DataFrame(data = tmp)
|
||
|
|
data = vectorizer.transform(df[0])
|
||
|
|
evaluate = linear.predict(data)
|
||
|
|
np.savetxt(out, evaluate, fmt='%f', delimiter='\n')
|
||
|
|
|
||
|
|
calculateResult("./dev-0/in.tsv", "./dev-0/out.tsv")
|
||
|
|
calculateResult("./dev-1/in.tsv", "./dev-1/out.tsv")
|
||
|
|
calculateResult("./test-A/in.tsv", "./test-A/out.tsv")
|