Как сделать прогноз из регрессионной модели?

Question

Я работаю через интро машинного обучения. В настоящее время я следил за классом, чтобы построить регрессионную модель для заработной платы, основанную на многолетнем опыте. Просто взглянув на регрессионную модель, я могу видеть, сколько зарабатывает кто-то, скажем, с 4-летним стажем, но мне интересно, есть ли функция или фрагмент кода, который может вернуть реальную прогнозируемую зарплату.

from sklearn.linear_model import LinearRegression
regressor = LinearRegression()
regressor.fit(X_train, y_train)


y_pred = regressor.predict(X_test)

#Training Set
plt.scatter(X_train, y_train, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title('Salary vs Experience (Training set)')
plt.xlabel('Years of Experience')
plt.ylabel('Salary')
plt.show()

#Test Results
plt.scatter(X_test, y_test, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title('Salary vs Experience (Training set)')
plt.xlabel('Years of Experience')
plt.ylabel('Salary')
plt.show()

Код выше - модель. Ниже приведен мой код для прогнозирования, но он не работает. Любая идея, как я могу адаптировать функцию для возврата функции на основе модели?

def PredictSalary(experience):
    experience = np.arange(6).reshape((3,2))
    regressor.predict(experience)

PredictSalary(6)

Мой код выше пытается превратить массив 1d в массив 2d, а затем использовать модель, чтобы сделать прогнозна зарплату за 6 лет.