Update main.py

This commit is contained in:
Iwona Christop 2022-05-08 23:44:48 +02:00
parent db6c2041ec
commit de955657c7
1 changed files with 2 additions and 1 deletions

View File

@ -32,7 +32,8 @@ if __name__ == '__main__':
train['Y'] = train.apply(lambda x: (float(x.Beginning) + float(x.End))/2, axis=1)
train = train.drop(columns=['Beginning', 'End', 'Title', 'Source'])
model = Pipeline([TfidfVectorizer(), LinearRegression()])
estimators = [('tfidf', TfidfVectorizer()), ('linearRegression', LinearRegression())]
model = Pipeline(estimators)
model.fit(train.X, train.Y)
# dev-0