File size: 3,828 Bytes
6a3d3b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
# train.py
import torch
import pandas as pd
import numpy as np
from torch import nn
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
class MangaDataset(Dataset):
def __init__(self, ratings_df):
self.users = torch.tensor(ratings_df['user_idx'].values, dtype=torch.long)
self.items = torch.tensor(ratings_df['manga_idx'].values, dtype=torch.long)
self.ratings = torch.tensor(ratings_df['rating'].values, dtype=torch.float)
def __len__(self):
return len(self.ratings)
def __getitem__(self, idx):
return self.users[idx], self.items[idx], self.ratings[idx]
class MangaRecommender(nn.Module):
def __init__(self, num_users, num_items, n_factors=50):
super().__init__()
self.user_factors = nn.Embedding(num_users, n_factors)
self.item_factors = nn.Embedding(num_items, n_factors)
# Initialize embeddings
nn.init.xavier_normal_(self.user_factors.weight)
nn.init.xavier_normal_(self.item_factors.weight)
def forward(self, user, item):
user_emb = self.user_factors(user)
item_emb = self.item_factors(item)
return (user_emb * item_emb).sum(1)
def predict(self, user_ids):
user_emb = self.user_factors(user_ids)
all_items = self.item_factors.weight
return torch.matmul(user_emb, all_items.t())
def train_model():
# Load your data
df = pd.read_csv('manga_ratings.csv')
# Create user and item mappings
user_mapping = {uid: idx for idx, uid in enumerate(df['user_id'].unique())}
manga_mapping = {mid: idx for idx, mid in enumerate(df['manga_id'].unique())}
# Convert ratings to numerical values
rating_map = {'like': 1.0, 'dislike': -1.0, None: 0.0}
# Prepare training data
df['user_idx'] = df['user_id'].map(user_mapping)
df['manga_idx'] = df['manga_id'].map(manga_mapping)
df['rating'] = df['like_status'].map(rating_map)
# Create train/val split
train_df, val_df = train_test_split(df, test_size=0.2)
# Create datasets
train_dataset = MangaDataset(train_df)
val_dataset = MangaDataset(val_df)
# Create dataloaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64)
# Initialize model
model = MangaRecommender(
num_users=len(user_mapping),
num_items=len(manga_mapping)
)
# Training setup
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# Training loop
num_epochs = 20
for epoch in range(num_epochs):
model.train()
total_loss = 0
for user, item, rating in train_loader:
optimizer.zero_grad()
pred = model(user, item)
loss = criterion(pred, rating)
loss.backward()
optimizer.step()
total_loss += loss.item()
# Validation
model.eval()
val_loss = 0
with torch.no_grad():
for user, item, rating in val_loader:
pred = model(user, item)
val_loss += criterion(pred, rating).item()
print(f'Epoch {epoch+1}/{num_epochs}')
print(f'Train Loss: {total_loss/len(train_loader):.4f}')
print(f'Val Loss: {val_loss/len(val_loader):.4f}')
# Save mappings and model
torch.save({
'model_state_dict': model.state_dict(),
'user_mapping': user_mapping,
'manga_mapping': manga_mapping
}, 'manga_recommender.pt')
if __name__ == '__main__':
train_model() |