Viraj45's picture
updated
fbbcc00 verified
# app.py
import torch
import pickle
import gradio as gr
from model import NCFModel
import numpy as np
# Load encoders and user_positive_items
with open('user_encoder.pkl', 'rb') as f:
user_encoder = pickle.load(f)
with open('item_encoder.pkl', 'rb') as f:
item_encoder = pickle.load(f)
with open('user_positive_items.pkl', 'rb') as f:
user_positive_items = pickle.load(f)
# Load the trained model
class NCFModelWrapper:
def __init__(self, model_path, num_users, num_items, embedding_size=50, device='cpu'):
self.device = torch.device(device)
self.model = NCFModel(num_users, num_items, embedding_size=embedding_size).to(self.device)
self.model.load_state_dict(torch.load(model_path, map_location=self.device))
self.model.eval()
def predict(self, user, item):
with torch.no_grad():
user = torch.tensor([user], dtype=torch.long).to(self.device)
item = torch.tensor([item], dtype=torch.long).to(self.device)
output = self.model(user, item)
score = torch.sigmoid(output).item()
return score
# Determine number of users and items from encoders
num_users = len(user_encoder.classes_)
num_items = len(item_encoder.classes_)
# Initialize the model
model = NCFModelWrapper(
model_path='best_ncf_model.pth',
num_users=num_users,
num_items=num_items,
embedding_size=50, # Ensure this matches your trained model
device='cpu' # Change to 'cuda' if GPU is available and desired
)
def recommend(user_id, num_recommendations=5):
try:
user = user_encoder.transform([user_id])[0]
except:
return f"User ID '{user_id}' not found."
# Get items the user has interacted with
pos_items = user_positive_items.get(user, set())
# Get all possible items
all_items = set(range(num_items))
# Candidate items are those not interacted with
candidate_items = list(all_items - pos_items)
# Predict scores for candidate items
scores = []
for item in candidate_items:
score = model.predict(user, item)
scores.append((item, score))
# Sort items based on score
scores.sort(key=lambda x: x[1], reverse=True)
# Get top N recommendations
top_items = scores[:num_recommendations]
recommendations = []
for item_id, score in top_items:
original_item_id = item_encoder.inverse_transform([item_id])[0]
recommendations.append(f"Item ID: {original_item_id} (Score: {score:.4f})")
return "\n".join(recommendations)
# Define Gradio interface
iface = gr.Interface(
fn=recommend,
inputs=[
gr.Textbox(lines=1, placeholder="Enter User ID", label="User ID"),
gr.Slider(minimum=1, maximum=20, step=1, value=5, label="Number of Recommendations")
],
outputs="text",
title="Neural Collaborative Filtering Recommendation System",
description="Enter a User ID to receive personalized item recommendations.",
examples=[
["user_1", 5],
["user_2", 10],
["user_3", 7]
]
)
if __name__ == "__main__":
iface.launch()