scripts/generate_data.py - master - katherinesiv/study_material_recommender

scripts/generate_data.py
214 строк7 KB
Сиваева Екатерина
update scripts/generate_data.py
26 дек 2025, 18:20
26 дек 2025, 18:205653956
#!/usr/bin/env python3
"""
Script to generate sample data for the Study Material Recommender system.
This script uses the data generator from the data directory.
"""

import argparse
from datetime import datetime
import os
import sys

# Fix import path before importing local modules
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

try:
    # Import the data generator module
    import data.sample_data_generator as sdg
except ImportError:
    print("Error: Could not import sample_data_generator module")
    print("Make sure you're running this script from the correct directory")
    sys.exit(1)


def main():
    """Generate sample data based on command line arguments"""
    parser = argparse.ArgumentParser(
        description="Generate sample data for Study Material Recommender"
    )

    parser.add_argument(
        "--output-dir",
        default="data",
        help="Directory to save generated CSV files"
    )

    parser.add_argument(
        "--n-materials",
        "--materials",
        type=int,
        default=100,
        help="Number of study materials to generate"
    )

    parser.add_argument(
        "--n-users",
        "--students",
        type=int,
        default=50,
        help="Number of users/students to generate"
    )

    parser.add_argument(
        "--rating-density",
        type=float,
        default=0.15,
        help="Density of ratings (0.0 to 1.0)"
    )

    parser.add_argument(
        "--generate-interactions",
        action="store_true",
        help="Generate additional interaction data"
    )

    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing files"
    )

    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducibility"
    )

    args = parser.parse_args()

    # Create output directory if it doesn't exist
    os.makedirs(args.output_dir, exist_ok=True)

    # Check if files exist
    materials_path = os.path.join(args.output_dir, "materials.csv")
    ratings_path = os.path.join(args.output_dir, "ratings.csv")

    if (
        os.path.exists(materials_path)
        and os.path.exists(ratings_path)
        and not args.overwrite
    ):
        print(
            f"Files already exist in {args.output_dir}. "
            "Use --overwrite to regenerate."
        )
        return

    print("=" * 60)
    print("Generating sample data for Study Material Recommender")
    print("=" * 60)
    print(f"Output directory: {args.output_dir}")
    print(f"Number of materials: {args.n_materials}")
    print(f"Number of users: {args.n_users}")
    print(f"Rating density: {args.rating_density}")
    print(f"Random seed: {args.seed}")
    print("=" * 60)

    # Initialize generator
    generator = sdg.SampleDataGenerator()

    # Generate materials
    print("\n1. Generating study materials...")
    materials_df = generator.generate_materials(args.n_materials)

    # Generate ratings
    print("\n2. Generating ratings...")
    ratings_df = generator.generate_ratings(
        n_users=args.n_users,
        n_materials=args.n_materials,
        rating_density=args.rating_density
    )

    # Save to CSV
    print("\n3. Saving data to CSV files...")
    materials_df.to_csv(materials_path, index=False)
    print(f"   Materials saved to: {materials_path}")

    ratings_df.to_csv(ratings_path, index=False)
    print(f"   Ratings saved to: {ratings_path}")

    # Generate interactions if requested
    if args.generate_interactions:
        print("\n4. Generating interaction data...")
        interactions_path = os.path.join(args.output_dir, "interactions.csv")
        interactions_df = generator.generate_interaction_data(n_interactions=2000)
        interactions_df.to_csv(interactions_path, index=False)
        print(f"   Interactions saved to: {interactions_path}")

    # Generate dataset statistics
    print("\n5. Generating dataset statistics...")

    try:
        import src.data_loader as dl

        data_loader = dl.DataLoader(materials_path, ratings_path)
        data = data_loader.preprocess_data()
        stats = data["stats"]

        stats_path = os.path.join(args.output_dir, "dataset_statistics.txt")
        with open(stats_path, "w") as f:
            f.write("=" * 60 + "\n")
            f.write("STUDY MATERIAL RECOMMENDER - DATASET STATISTICS\n")
            f.write("=" * 60 + "\n\n")

            f.write("DATASET OVERVIEW:\n")
            f.write("-" * 40 + "\n")
            f.write(f"Generated on: {datetime.now()}\n")
            f.write(f"Number of materials: {stats['num_materials']}\n")
            f.write(f"Number of users: {stats['num_users']}\n")
            f.write(f"Number of ratings: {len(ratings_df)}\n")
            f.write(f"Average rating: {stats['avg_rating']:.2f}\n")
            f.write(f"Rating sparsity: {stats['rating_sparsity']:.3f}\n\n")

            f.write("MATERIALS BY SUBJECT:\n")
            f.write("-" * 40 + "\n")
            subject_counts = materials_df["subject"].value_counts()
            for subject, count in subject_counts.items():
                f.write(f"{subject}: {count} materials\n")
            f.write("\n")

            f.write("MATERIALS BY DIFFICULTY:\n")
            f.write("-" * 40 + "\n")
            difficulty_counts = materials_df["difficulty"].value_counts()
            for difficulty, count in difficulty_counts.items():
                f.write(f"{difficulty}: {count} materials\n")
            f.write("\n")

            f.write("RATING DISTRIBUTION:\n")
            f.write("-" * 40 + "\n")
            rating_dist = ratings_df["rating"].value_counts().sort_index()
            for rating, count in rating_dist.items():
                percentage = (count / len(ratings_df)) * 100
                f.write(f"Rating {rating}: {count} ratings ({percentage:.1f}%)\n")

        print(f"   Statistics saved to: {stats_path}")

    except Exception as e:
        print(f"   Warning: Could not generate detailed statistics: {e}")

    print("\n" + "=" * 60)
    print("DATA GENERATION COMPLETE!")
    print("=" * 60)

    # Quick verification
    print("\nVERIFICATION:")
    materials_size = os.path.getsize(materials_path) / 1024
    ratings_size = os.path.getsize(ratings_path) / 1024
    print(f"- Materials file size: {materials_size:.1f} KB")
    print(f"- Ratings file size: {ratings_size:.1f} KB")

    print("\nLOADED DATA SUMMARY:")
    print(f"- Materials shape: {materials_df.shape}")
    print(f"- Ratings shape: {ratings_df.shape}")
    print(f"- Unique users in ratings: {ratings_df['user_id'].nunique()}")
    print(
        "- Unique materials in ratings: "
        f"{ratings_df['material_id'].nunique()}"
    )

    print("\nThe dataset is ready for use with the Study Material Recommender!")


if __name__ == "__main__":
    main()