price-builder-backend/services/modelBuilder.py
2024-10-13 16:53:07 +04:00

144 lines
5.7 KiB
Python
Raw Permalink 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.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import joblib
import re
# Шаг 1: Загрузка данных
df = pd.read_csv('../laptops.csv')
# Шаг 2: Проверка и очистка имен столбцов
print("Имена столбцов до очистки:")
print(df.columns.tolist())
# Приведение имен столбцов к нижнему регистру и удаление пробелов
df.columns = df.columns.str.strip().str.lower()
print("\nИмена столбцов после очистки:")
print(df.columns.tolist())
# Шаг 3: Переименование столбцов (если необходимо)
df = df.rename(columns={
'processor': 'processor',
'ram': 'ram',
'os': 'os',
'ssd': 'ssd',
'display': 'display',
'price': 'price'
# Другие столбцы можно оставить без изменений или переименовать по необходимости
})
# Шаг 4: Проверка наличия необходимых столбцов
required_columns = ['processor', 'ram', 'os', 'ssd', 'display', '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Все необходимые столбцы присутствуют.")
# Шаг 5: Удаление строк с пропущенными значениями
df = df.dropna(subset=required_columns)
print(f"\nКоличество строк после удаления пропусков: {df.shape[0]}")
# Шаг 6: Очистка и преобразование колонок
# Функция для очистки числовых колонок, содержащих символы
def clean_numeric_column(column, remove_chars=['', ',', ' ']):
for char in remove_chars:
column = column.str.replace(char, '', regex=False)
return pd.to_numeric(column, errors='coerce')
# Очистка колонки 'price'
df['price'] = clean_numeric_column(df['price'])
# Очистка колонки 'ram' (например, '16 GB DDR4 RAM' -> 16)
def extract_numeric_ram(ram_str):
match = re.search(r'(\d+)', ram_str)
if match:
return int(match.group(1))
else:
return None
df['ram'] = df['ram'].apply(extract_numeric_ram)
# Очистка колонки 'ssd' (например, '512 GB SSD' -> 512)
def extract_numeric_ssd(ssd_str):
match = re.search(r'(\d+)', ssd_str)
if match:
return int(match.group(1))
else:
return None
df['ssd'] = df['ssd'].apply(extract_numeric_ssd)
# Очистка колонки 'display' (убираем лишние символы, если есть)
def clean_display(display_str):
match = re.search(r'([\d.]+)', display_str)
if match:
return float(match.group(1))
else:
return None
df['display'] = df['display'].apply(clean_display)
# Проверка на пропущенные значения после очистки
df = df.dropna(subset=['price', 'ram', 'ssd', 'display'])
print(f"\nКоличество строк после очистки числовых колонок: {df.shape[0]}")
# Шаг 7: Выбор необходимых столбцов
df = df[required_columns]
print("\nПример данных после предобработки:")
print(df.head())
# Шаг 8: Преобразование категориальных переменных с помощью One-Hot Encoding
df = pd.get_dummies(df, columns=['processor', 'os'], drop_first=True)
print("\nИмена колонок после One-Hot Encoding:")
print(df.columns.tolist())
# Шаг 9: Разделение данных на обучающую и тестовую выборки
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}")
# Шаг 10: Обучение модели
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
# Обучение моделей
gbr = GradientBoostingRegressor(n_estimators=100, random_state=42)
gbr.fit(X_train, y_train)
lr = LinearRegression()
lr.fit(X_train, y_train)
print("\nМодели успешно обучена.")
# Оценка моделей
models = {'Random Forest': model, 'Gradient Boosting': gbr, 'Linear Regression': lr}
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}")
# Шаг 12: Сохранение модели
joblib.dump(model, '../laptop_price_model.pkl')
print("\nМодель сохранена как 'laptop_price_model.pkl'.")
# Дополнительно: Сохранение колонок, полученных после One-Hot Encoding, для использования в бэкенде
feature_columns = X.columns.tolist()
joblib.dump(feature_columns, '../feature_columns.pkl')
print("Сохранены названия признаков в 'feature_columns.pkl'.")