| """ |
| Synthetic MovieLens-style dataset with realistic user-item interactions. |
| Generates data using a latent factor model if no data file exists. |
| """ |
| import os |
| import numpy as np |
| import pandas as pd |
|
|
|
|
| class MovieLensDataset: |
| GENRES = [ |
| 'Action', 'Comedy', 'Drama', 'Horror', 'Sci-Fi', |
| 'Romance', 'Thriller', 'Animation', 'Documentary', 'Fantasy', |
| ] |
|
|
| def __init__(self, data_dir='data', n_users=500, n_items=200): |
| self.data_dir = data_dir |
| self.n_users = n_users |
| self.n_items = n_items |
| os.makedirs(data_dir, exist_ok=True) |
|
|
| ratings_path = os.path.join(data_dir, 'ratings.csv') |
| items_path = os.path.join(data_dir, 'items.csv') |
|
|
| if os.path.exists(ratings_path) and os.path.exists(items_path): |
| self.ratings_df = pd.read_csv(ratings_path) |
| self.items_df = pd.read_csv(items_path) |
| else: |
| self._generate_synthetic() |
|
|
| self._preprocess() |
|
|
| |
| def _generate_synthetic(self): |
| np.random.seed(42) |
| n_factors = 20 |
| user_factors = np.random.randn(self.n_users, n_factors) |
| item_factors = np.random.randn(self.n_items, n_factors) |
|
|
| interactions = [] |
| for user_id in range(self.n_users): |
| n_ratings = np.random.randint(20, 80) |
| scores = user_factors[user_id] @ item_factors.T |
| probs = np.exp(scores - scores.max()) |
| probs /= probs.sum() |
| n_sample = min(n_ratings, self.n_items) |
| rated_items = np.random.choice(self.n_items, size=n_sample, replace=False, p=probs) |
| for item_id in rated_items: |
| score = user_factors[user_id] @ item_factors[item_id] |
| rating = int(np.clip(round((score + 2) / 4 * 4 + 1), 1, 5)) |
| timestamp = np.random.randint(1_000_000, 2_000_000) |
| interactions.append((user_id, item_id, rating, timestamp)) |
|
|
| self.ratings_df = pd.DataFrame( |
| interactions, columns=['user_id', 'item_id', 'rating', 'timestamp'] |
| ) |
| self.ratings_df = self.ratings_df.sort_values('timestamp').reset_index(drop=True) |
|
|
| items = [] |
| for item_id in range(self.n_items): |
| n_genres = np.random.randint(1, 4) |
| item_genres = np.random.choice(self.GENRES, size=n_genres, replace=False) |
| items.append({ |
| 'item_id': item_id, |
| 'title': f'Movie {item_id + 1}', |
| 'genres': '|'.join(item_genres), |
| }) |
| self.items_df = pd.DataFrame(items) |
|
|
| self.ratings_df.to_csv(os.path.join(self.data_dir, 'ratings.csv'), index=False) |
| self.items_df.to_csv(os.path.join(self.data_dir, 'items.csv'), index=False) |
|
|
| |
| def _preprocess(self): |
| df = self.ratings_df.copy() |
| df['reward'] = (df['rating'] >= 4).astype(int) |
|
|
| self.user_sequences: dict = {} |
| for user_id, group in df.groupby('user_id'): |
| group_sorted = group.sort_values('timestamp') |
| self.user_sequences[int(user_id)] = list(zip( |
| group_sorted['item_id'].values.tolist(), |
| group_sorted['reward'].values.tolist(), |
| )) |
|
|
| train_rows, test_rows = [], [] |
| for user_id, seq in self.user_sequences.items(): |
| split = max(1, int(len(seq) * 0.8)) |
| for i, (item_id, reward) in enumerate(seq): |
| row = {'user_id': user_id, 'item_id': item_id, 'reward': reward} |
| (train_rows if i < split else test_rows).append(row) |
|
|
| self.train_df = pd.DataFrame(train_rows) |
| self.test_df = pd.DataFrame(test_rows) |
| self.n_users_actual = int(df['user_id'].nunique()) |
| self.n_items_actual = int(df['item_id'].nunique()) |
|
|
| |
| def get_user_history(self, user_id: int, max_len: int = 20): |
| return self.user_sequences.get(user_id, [])[-max_len:] |
|
|
| def get_item_info(self, item_id: int) -> dict: |
| row = self.items_df[self.items_df['item_id'] == item_id] |
| if len(row) == 0: |
| return {'item_id': item_id, 'title': f'Movie {item_id + 1}', 'genres': 'Unknown'} |
| r = row.iloc[0] |
| return {'item_id': int(item_id), 'title': str(r['title']), 'genres': str(r['genres'])} |
|
|
