src/model.py - main - NeuroEdTech/AI-marafon

src/model.py

22 строки672 B

Oleg

Initial commit for AI-marafon NLP ptoject

24 янв 2026, 23:47

24 янв 2026, 23:477601e61

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import joblib
from pathlib import Path

def train_model(X_train, y_train, n_estimators=200, random_state=42):
    model = RandomForestClassifier(
        n_estimators=n_estimators,
        random_state=random_state,
        n_jobs=-1
    )
    model.fit(X_train, y_train)
    return model

def save_model(model, path: str = "models/model.pkl"):
    Path("models").mkdir(exist_ok=True)
    joblib.dump(model, path)

def train_and_save(X_train, y_train, model_path="models/model.pkl"):
    model = train_model(X_train, y_train)
    save_model(model, model_path)
    return model