src/main.py - master - katherinesiv/study_material_recommender

src/main.py
372 строки12 KB
Сиваева Екатерина
update src/main.py
25 дек 2025, 21:48
25 дек 2025, 21:48496a231
"""
Main script for training and evaluating recommendation models
"""

import argparse
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Tuple

import numpy as np
import pandas as pd

from .api import app
from .collaborative_filtering import CollaborativeFiltering
from .content_based_filtering import ContentBasedFiltering
from .data_loader import DataLoader
from .evaluation import RecommendationEvaluator
from .hybrid_model import HybridRecommender


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def train_models(data_path: str = "data", output_path: str = "models") -> Dict[str, Any]:
    """
    Train all recommendation models and save results.

    Args:
        data_path: Path to data directory
        output_path: Path to save trained models and metrics

    Returns:
        Dictionary of evaluation metrics
    """
    logger.info("Starting model training...")

    # Create output directory
    Path(output_path).mkdir(parents=True, exist_ok=True)

    # Load and preprocess data
    data_loader = DataLoader(
        materials_path=f"{data_path}/materials.csv",
        ratings_path=f"{data_path}/ratings.csv"
    )

    data = data_loader.preprocess_data()
    logger.info(f"Data loaded: {data['stats']}")

    # Split data into train and test
    train_data, test_data = split_data(data["ratings_df"])

    # Train and evaluate Collaborative Filtering
    logger.info("Training Collaborative Filtering model...")
    cf_model = train_collaborative_filtering(train_data, data["user_item_matrix"])
    cf_metrics = evaluate_model(cf_model, test_data, "collaborative")

    # Train and evaluate Content-Based Filtering
    logger.info("Training Content-Based Filtering model...")
    cb_model = train_content_based_filtering(data["materials_processed"])
    cb_metrics = evaluate_content_based(cb_model, test_data, data)

    # Train and evaluate Hybrid Model
    logger.info("Training Hybrid model...")
    hybrid_model = train_hybrid_model(data)
    hybrid_metrics = evaluate_hybrid(hybrid_model, test_data, data)

    # Save metrics
    metrics = {
        "timestamp": datetime.now().isoformat(),
        "data_statistics": data["stats"],
        "collaborative_filtering": cf_metrics,
        "content_based": cb_metrics,
        "hybrid": hybrid_metrics,
        "best_model": determine_best_model(cf_metrics, cb_metrics, hybrid_metrics)
    }

    metrics_file = f"{output_path}/metrics_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    with open(metrics_file, "w") as f:
        json.dump(metrics, f, indent=2)

    logger.info(f"Metrics saved to {metrics_file}")

    # Save model artifacts (in production, you'd use pickle/joblib)
    save_model_artifacts(output_path, {
        "cf_model": cf_model,
        "cb_model": cb_model,
        "hybrid_model": hybrid_model,
        "data_loader": data_loader
    })

    logger.info("Model training completed successfully!")
    print("\n=== Training Summary ===")
    print(f"Best model: {metrics['best_model']['model']}")
    print(f"Best RMSE: {metrics['best_model']['rmse']:.3f}")
    print(f"Best Precision@10: {metrics['best_model']['precision@10']:.3f}")

    return metrics


def split_data(ratings_df: pd.DataFrame, test_size: float = 0.2) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Split ratings data into train and test sets.

    Args:
        ratings_df: Ratings DataFrame
        test_size: Proportion of data to use for testing

    Returns:
        Tuple of (train_data, test_data)
    """
    # Simple time-based split (if timestamp exists)
    if "timestamp" in ratings_df.columns:
        ratings_df = ratings_df.sort_values("timestamp")
        split_idx = int(len(ratings_df) * (1 - test_size))
        train_data = ratings_df.iloc[:split_idx]
        test_data = ratings_df.iloc[split_idx:]
    else:
        # Random split
        mask = np.random.rand(len(ratings_df)) < (1 - test_size)
        train_data = ratings_df[mask]
        test_data = ratings_df[~mask]

    logger.info(f"Train size: {len(train_data)}, Test size: {len(test_data)}")
    return train_data, test_data


def train_collaborative_filtering(
    train_data: pd.DataFrame, user_item_matrix: pd.DataFrame
) -> CollaborativeFiltering:
    """Train collaborative filtering model"""
    # Create train user-item matrix
    train_matrix = train_data.pivot_table(
        index="user_id",
        columns="material_id",
        values="rating",
        fill_value=0
    )

    # Align with full matrix columns
    train_matrix = train_matrix.reindex(columns=user_item_matrix.columns, fill_value=0)

    model = CollaborativeFiltering(k_neighbors=15)
    model.fit(train_matrix)
    return model


def train_content_based_filtering(materials_processed: pd.DataFrame) -> ContentBasedFiltering:
    """Train content-based filtering model"""
    model = ContentBasedFiltering()
    model.fit(materials_processed)
    return model


def train_hybrid_model(data: Dict[str, Any]) -> HybridRecommender:
    """Train hybrid recommendation model"""
    model = HybridRecommender()
    model.fit(data)
    return model


def evaluate_model(
    model, test_data: pd.DataFrame, model_type: str
) -> Dict[str, Any]:
    """
    Evaluate a recommendation model.

    Args:
        model: Trained model
        test_data: Test ratings
        model_type: Type of model

    Returns:
        Dictionary of evaluation metrics
    """
    evaluator = RecommendationEvaluator(test_data)

    # Generate predictions and recommendations
    predictions = []
    recommendations = {}

    # Sample evaluation for demonstration
    sample_users = test_data["user_id"].unique()[:50]

    for user_id in sample_users:
        try:
            if model_type == "collaborative":
                user_recs = model.recommend(user_id, n_recommendations=10)
                rec_ids = [r["material_id"] for r in user_recs]
            else:
                continue

            recommendations[user_id] = rec_ids

            # Generate sample predictions
            user_test_items = test_data[test_data["user_id"] == user_id]
            for _, row in user_test_items.iterrows():
                if model_type == "collaborative":
                    pred = model.predict_rating(user_id, row["material_id"])
                    predictions.append({
                        "user_id": user_id,
                        "material_id": row["material_id"],
                        "predicted_rating": pred
                    })

        except Exception as e:
            logger.debug(f"Error evaluating user {user_id}: {e}")
            continue

    predictions_df = pd.DataFrame(predictions) if predictions else None

    model_predictions = {
        "predictions": predictions_df,
        "recommendations": recommendations
    }

    metrics = evaluator.evaluate_model(model_predictions)
    metrics["model_type"] = model_type

    return metrics


def evaluate_content_based(
    model: ContentBasedFiltering, test_data: pd.DataFrame, data: Dict[str, Any]
) -> Dict[str, Any]:
    """Evaluate content-based filtering model"""
    evaluator = RecommendationEvaluator(test_data)
    recommendations = {}

    # Sample evaluation
    sample_users = test_data["user_id"].unique()[:50]

    for user_id in sample_users:
        user_history = data["ratings_df"][data["ratings_df"]["user_id"] == user_id]
        recs = model.recommend_for_user(user_history, n_recommendations=10)
        rec_ids = [r["material_id"] for r in recs]
        recommendations[user_id] = rec_ids

    model_predictions = {"recommendations": recommendations}
    metrics = evaluator.evaluate_model(model_predictions)
    metrics["model_type"] = "content_based"

    return metrics


def evaluate_hybrid(
    model: HybridRecommender, test_data: pd.DataFrame, data: Dict[str, Any]
) -> Dict[str, Any]:
    """Evaluate hybrid model"""
    evaluator = RecommendationEvaluator(test_data)
    predictions = []
    recommendations = {}

    sample_users = test_data["user_id"].unique()[:50]

    for user_id in sample_users:
        try:
            user_history = data["ratings_df"][data["ratings_df"]["user_id"] == user_id]
            recs = model.recommend(user_id, n_recommendations=10, user_history=user_history)
            rec_ids = [r["material_id"] for r in recs]
            recommendations[user_id] = rec_ids

            # Generate sample predictions (simplified)
            user_test_items = test_data[test_data["user_id"] == user_id]
            for _, row in user_test_items.head(5).iterrows():
                predictions.append({
                    "user_id": user_id,
                    "material_id": row["material_id"],
                    "predicted_rating": 4.0
                })

        except Exception as e:
            logger.debug(f"Error evaluating user {user_id}: {e}")
            continue

    predictions_df = pd.DataFrame(predictions) if predictions else None

    model_predictions = {
        "predictions": predictions_df,
        "recommendations": recommendations
    }

    metrics = evaluator.evaluate_model(model_predictions)
    metrics["model_type"] = "hybrid"

    return metrics


def determine_best_model(
    cf_metrics: Dict[str, Any], cb_metrics: Dict[str, Any], hybrid_metrics: Dict[str, Any]
) -> Dict[str, Any]:
    """Determine the best performing model based on multiple metrics"""
    models = {
        "collaborative": cf_metrics,
        "content_based": cb_metrics,
        "hybrid": hybrid_metrics
    }

    # Simple scoring: lower RMSE and higher precision are better
    scores = {}
    for name, metrics in models.items():
        score = 0
        if "rmse" in metrics:
            score += 1 / (metrics["rmse"] + 0.001)
        if "precision@10" in metrics:
            score += metrics["precision@10"] * 10

        scores[name] = score

    best_model = max(scores.items(), key=lambda x: x[1])

    return {
        "model": best_model[0],
        "score": best_model[1],
        "rmse": models[best_model[0]].get("rmse", "N/A"),
        "precision@10": models[best_model[0]].get("precision@10", "N/A")
    }


def save_model_artifacts(output_path: str, artifacts: Dict[str, Any]) -> None:
    """Save model artifacts to disk"""
    logger.info(f"Saving model artifacts to {output_path}")

    # Save metadata about the models
    metadata = {
        "timestamp": datetime.now().isoformat(),
        "models": list(artifacts.keys()),
        "notes": "Model artifacts would be saved here in production"
    }

    with open(f"{output_path}/metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)


def run_api_server(host: str = "0.0.0.0", port: int = 8000) -> None:
    """Run the FastAPI server"""
    import uvicorn
    logger.info(f"Starting API server on {host}:{port}")
    uvicorn.run(app, host=host, port=port)


def main() -> None:
    """Main entry point"""
    parser = argparse.ArgumentParser(description="Study Material Recommender")
    parser.add_argument("--train", action="store_true", help="Train models")
    parser.add_argument("--serve", action="store_true", help="Start API server")
    parser.add_argument("--data-path", default="data", help="Path to data directory")
    parser.add_argument("--output-path", default="models", help="Path to save models")
    parser.add_argument("--host", default="0.0.0.0", help="API server host")
    parser.add_argument("--port", type=int, default=8000, help="API server port")

    args = parser.parse_args()

    if args.train:
        train_models(args.data_path, args.output_path)

    if args.serve:
        run_api_server(args.host, args.port)

    if not args.train and not args.serve:
        # Default: train models
        train_models(args.data_path, args.output_path)


if __name__ == "__main__":
    main()