book-recommender / model.py
nirajandhakal's picture
initial commit
350eabd verified
raw
history blame
No virus
2.68 kB
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Embedding, Flatten, concatenate, Dense
from tensorflow.keras.optimizers import Adam
# Load datasets
books = pd.read_csv("../data/dataset/books.csv")
ratings = pd.read_csv("../data/dataset/ratings.csv")
# Preprocess data
user_encoder = LabelEncoder()
book_encoder = LabelEncoder()
ratings["user_id"] = user_encoder.fit_transform(ratings["user_id"])
ratings["book_id"] = book_encoder.fit_transform(ratings["book_id"])
# Split the data into training and testing sets
train, test = train_test_split(ratings, test_size=0.2, random_state=42)
# Define the neural network model
def build_model(num_users, num_books, embedding_size=50):
"""
Build a recommendation model.
Args:
num_users (int): The number of users in the dataset.
num_books (int): The number of books in the dataset.
embedding_size (int, optional): The size of the embedding vectors. Defaults to 50.
Returns:
keras.Model: The compiled recommendation model.
"""
user_input = Input(shape=(1,))
book_input = Input(shape=(1,))
user_embedding = Embedding(input_dim=num_users, output_dim=embedding_size)(
user_input
)
book_embedding = Embedding(input_dim=num_books, output_dim=embedding_size)(
book_input
)
user_flat = Flatten()(user_embedding)
book_flat = Flatten()(book_embedding)
merged = concatenate([user_flat, book_flat])
dense1 = Dense(128, activation="relu")(merged)
output = Dense(1)(dense1)
model = Model(inputs=[user_input, book_input], outputs=output)
model.compile(loss="mean_squared_error", optimizer=Adam(learning_rate=0.001))
return model
# Train the model
model = build_model(
num_users=len(ratings["user_id"].unique()),
num_books=len(ratings["book_id"].unique()),
)
history = model.fit(
[train["user_id"], train["book_id"]],
train["rating"],
epochs=5,
batch_size=128,
validation_split=0.1,
)
# Plot training and validation loss
plt.figure(figsize=(12, 6))
plt.plot(history.history["loss"], label="Training Loss")
plt.plot(history.history["val_loss"], label="Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()
# Save the model
model.save("recommendation_model.h5")
# Evaluate the model
test_loss = model.evaluate([test["user_id"], test["book_id"]], test["rating"])
print(f"Test Loss: {test_loss}")