forked from kubapok/retroc2
56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.linear_model import LinearRegression
|
|
import pickle
|
|
|
|
filename = 'regressor.sav'
|
|
vec_file = 'vectorizer.pickle'
|
|
regressor = LinearRegression()
|
|
# regressor = pickle.load(open(filename, 'rb'))
|
|
vectorizer = TfidfVectorizer()
|
|
# vectorizer = pickle.load(open(vec_file, 'rb'))
|
|
|
|
|
|
def train():
|
|
with open('train/train.tsv') as f:
|
|
docs = [line.rstrip() for line in f]
|
|
docs_preprocessed = []
|
|
y = []
|
|
for doc in docs:
|
|
row = doc.split('\t')
|
|
start_date = row[0]
|
|
end_date = row[1]
|
|
end_date = end_date.split(' ')
|
|
if len(end_date) > 1:
|
|
row.insert(4, end_date[1])
|
|
end_date = end_date[0]
|
|
doc = row[4:5][0]
|
|
docs_preprocessed.append(doc)
|
|
y.append((float(start_date) + float(end_date))/2)
|
|
y = [float(value) for value in y]
|
|
print('Fitting vectorizer...')
|
|
x = vectorizer.fit_transform(docs_preprocessed)
|
|
pickle.dump(vectorizer, open(vec_file, 'wb'))
|
|
print('DONE!')
|
|
print('Fitting regressor...')
|
|
regressor.fit(x, y)
|
|
pickle.dump(regressor, open(filename, 'wb'))
|
|
print('DONE!')
|
|
|
|
|
|
def classify(path):
|
|
print("Predicting for", path)
|
|
with open(path + 'in.tsv') as f:
|
|
docs = [line.rstrip() for line in f]
|
|
test_x = vectorizer.transform(docs)
|
|
predictions = regressor.predict(test_x)
|
|
with open(path + 'out.tsv', 'w') as file:
|
|
for prediction in predictions:
|
|
file.write("%f\n" % prediction)
|
|
|
|
|
|
train()
|
|
classify('dev-0/')
|
|
classify('dev-1/')
|
|
classify('test-A/')
|
|
|