price-builder-backend/services/ml/modelBuilder.py

124 lines
5.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import matplotlib.pyplot as plt
import joblib
import numpy as np
# Шаг 1: Загрузка данных
df = pd.read_csv('../../datasets/synthetic_laptops.csv') # Убедитесь, что путь к файлу правильный
# Шаг 2: Проверка и очистка имен столбцов
print("Имена столбцов до очистки:")
print(df.columns.tolist())
# Приведение имен столбцов к нижнему регистру и удаление пробелов
df.columns = df.columns.str.strip().str.lower()
print("\nИмена столбцов после очистки:")
print(df.columns.tolist())
# Шаг 3: Проверка наличия необходимых столбцов
required_columns = [
'brand', 'processor', 'ram', 'os', 'ssd', 'display',
'gpu', 'weight', 'battery_size', 'release_year', 'display_type', 'price'
]
missing_columns = [col for col in required_columns if col not in df.columns]
if missing_columns:
print(f"\nОтсутствуют следующие столбцы: {missing_columns}")
raise Exception(f"Отсутствуют столбцы: {missing_columns}")
else:
print("\nВсе необходимые столбцы присутствуют.")
# Шаг 4: Удаление строк с пропущенными значениями
df = df.dropna(subset=required_columns)
print(f"\nКоличество строк после удаления пропусков: {df.shape[0]}")
# Шаг 5: Очистка и преобразование колонок
# Функция для очистки числовых колонок, если они строковые
def clean_numeric_column(column, remove_chars=['', ',', ' ']):
if column.dtype == object:
for char in remove_chars:
column = column.str.replace(char, '', regex=False)
return pd.to_numeric(column, errors='coerce')
else:
return column
# Очистка числовых колонок (исключая 'price', если уже числовая)
numerical_columns = ['ram', 'ssd', 'display', 'weight', 'battery_size', 'release_year']
for col in numerical_columns:
df[col] = clean_numeric_column(df[col])
# Проверка на пропущенные значения после очистки
df = df.dropna(subset=['price'] + numerical_columns)
print(f"\nКоличество строк после очистки числовых колонок: {df.shape[0]}")
# Шаг 6: Выбор необходимых столбцов (все уже включены)
# df уже содержит все необходимые столбцы
print("\nПример данных после предобработки:")
print(df.head())
# Шаг 7: Преобразование категориальных переменных с помощью One-Hot Encoding
categorical_features = ['brand', 'processor', 'os', 'gpu', 'display_type']
df = pd.get_dummies(df, columns=categorical_features, drop_first=True)
print("\nИмена колонок после One-Hot Encoding:")
print(df.columns.tolist())
# Шаг 8: Разделение данных на обучающую и тестовую выборки
X = df.drop('price', axis=1)
y = df['price']
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
print(f"\nРазмер обучающей выборки: {X_train.shape}")
print(f"Размер тестовой выборки: {X_test.shape}")
# Шаг 9: Обучение моделей
model_rf = RandomForestRegressor(n_estimators=100, random_state=42)
model_rf.fit(X_train, y_train)
print("\nМодели успешно обучены.")
# Шаг 10: Оценка моделей
models = {
'Random Forest': model_rf,
}
for name, mdl in models.items():
y_pred = mdl.predict(X_test)
mae = mean_absolute_error(y_test, y_pred)
rmse = mean_squared_error(y_test, y_pred, squared=False)
r2 = r2_score(y_test, y_pred)
print(f"{name} - MAE: {mae}, RMSE: {rmse}, R²: {r2}")
# Шаг 11: Сохранение модели и списка признаков
joblib.dump(model_rf, 'laptop_price_model.pkl')
print("\nМодель Random Forest сохранена как 'laptop_price_model.pkl'.")
feature_columns = X.columns.tolist()
joblib.dump(feature_columns, 'feature_columns.pkl')
print("Сохранены названия признаков в 'feature_columns.pkl'.")
# Получение важности признаков
importances = model_rf.feature_importances_
indices = np.argsort(importances)[::-1]
# Вывод наиболее важных признаков
print("Важность признаков:")
for f in range(X_train.shape[1]):
print(f"{f + 1}. {feature_columns[indices[f]]} ({importances[indices[f]]})")
# Визуализация важности признаков
plt.figure(figsize=(12, 8))
plt.title("Важность признаков (Random Forest)")
plt.bar(range(X_train.shape[1]), importances[indices], align='center')
plt.xticks(range(X_train.shape[1]), [feature_columns[i] for i in indices], rotation=90)
plt.tight_layout()
plt.show()