src/model_training.py - master - Dinky6/Assessment-Question-Difficulty-Classifier

src/model_training.py
99 строк4 KB
Dinky6
create src/model_training.py
28 дек 2025, 17:14
28 дек 2025, 17:14dab6901
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.naive_bayes import GaussianNB
from sklearn.model_selection import cross_val_score, GridSearchCV
import pandas as pd
import numpy as np
import joblib
import json
import os

class DifficultyClassifier:
    """Классификатор сложности вопросов"""
    
    def __init__(self, model_type='random_forest'):
        self.model_type = model_type
        self.model = None
        self.feature_extractor = None
        self.scaler = None
        self.encoder = None
        
        if model_type == 'random_forest':
            self.model = RandomForestClassifier(
                n_estimators=100,
                max_depth=10,
                random_state=42,
                class_weight='balanced'
            )
        elif model_type == 'svm':
            self.model = SVC(
                kernel='rbf',
                C=1.0,
                probability=True,
                random_state=42,
                class_weight='balanced'
            )
        elif model_type == 'naive_bayes':
            self.model = GaussianNB()
        else:
            raise ValueError(f"Unknown model type: {model_type}")
    
    def train(self, X_train, y_train, cv_folds=5):
        """Обучение модели с кросс-валидацией"""
        cv_scores = cross_val_score(self.model, X_train, y_train, cv=cv_folds)
        print(f"Cross-validation scores: {cv_scores}")
        print(f"Mean CV accuracy: {cv_scores.mean():.3f} (+/- {cv_scores.std() * 2:.3f})")
        
        self.model.fit(X_train, y_train)
        return self
    
    def predict(self, X):
        """Предсказание сложности"""
        return self.model.predict(X)
    
    def predict_proba(self, X):
        """Вероятности предсказаний"""
        return self.model.predict_proba(X)
    
    def save_model(self, path='models/question_classifier.pkl'):
        """Сохранение модели"""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        joblib.dump({
            'model': self.model,
            'model_type': self.model_type,
            'feature_extractor': self.feature_extractor,
            'scaler': self.scaler,
            'encoder': self.encoder
        }, path)
        print(f"Model saved to {path}")
    
    @classmethod
    def load_model(cls, path='models/question_classifier.pkl'):
        """Загрузка модели"""
        data = joblib.load(path)
        classifier = cls(data['model_type'])
        classifier.model = data['model']
        classifier.feature_extractor = data['feature_extractor']
        classifier.scaler = data['scaler']
        classifier.encoder = data['encoder']
        return classifier
    
    def hyperparameter_tuning(self, X_train, y_train):
        """Оптимизация гиперпараметров"""
        if self.model_type == 'random_forest':
            param_grid = {
                'n_estimators': [50, 100, 200],
                'max_depth': [5, 10, 20, None],
                'min_samples_split': [2, 5, 10],
                'min_samples_leaf': [1, 2, 4]
            }
            grid_search = GridSearchCV(
                self.model, param_grid, cv=5, 
                scoring='accuracy', n_jobs=-1
            )
            grid_search.fit(X_train, y_train)
            self.model = grid_search.best_estimator_
            print(f"Best parameters: {grid_search.best_params_}")
            print(f"Best CV accuracy: {grid_search.best_score_:.3f}")
        
        return self