:param data: 训练数据 :param education: 学历 :return: 模型得分，10年工资预测
(data, education)
| 9 | from sklearn.linear_model import LinearRegression |
| 10 | |
| 11 | def predict(data, education): |
| 12 | """ |
| 13 | :param data: 训练数据 |
| 14 | :param education: 学历 |
| 15 | :return: 模型得分，10年工资预测 |
| 16 | """ |
| 17 | train = data[data['education'] == education].to_numpy() |
| 18 | x = train[:, 1:2] |
| 19 | y = train[:, 2] |
| 20 | |
| 21 | # model 训练 |
| 22 | model = LinearRegression() |
| 23 | model.fit(x, y) |
| 24 | |
| 25 | # model 预测 |
| 26 | X = [[i] for i in range(11)] |
| 27 | return model.score(x, y), model.predict(X) |
| 28 | |
| 29 | education_list = ['小学', '初中', '中专', '高中', '大专', '本科', '硕士', '博士'] |
| 30 | data = pd.read_csv('train.csv') |