clfegg commited on
Commit
1e87c85
1 Parent(s): 8483e67

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +15 -1
handler.py CHANGED
@@ -2,8 +2,22 @@
2
  from typing import Dict, List, Any
3
  import pickle
4
  import os
5
- from model import ContentBasedRecommender
6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7
  class EndpointHandler:
8
  def __init__(self, path=""):
9
  model_path = os.path.join(path, "model.pkl")
 
2
  from typing import Dict, List, Any
3
  import pickle
4
  import os
5
+ import __main__
6
 
7
+
8
+ class ContentBasedRecommender:
9
+ def __init__(self, train_data):
10
+ self.train_data = train_data
11
+
12
+ def predict(self, user_id, k=10):
13
+ user_books = set(self.train_data[self.train_data['user_id'] == user_id]['book_id'])
14
+ similar_books = set().union(*(self.train_data[self.train_data['book_id'] == book_id]['similar_books'].iloc[0] for book_id in user_books))
15
+ recommended_books = list(similar_books - user_books)
16
+
17
+ return np.random.choice(recommended_books, size=min(k, len(recommended_books)), replace=False).tolist()
18
+
19
+
20
+ __main__.ContentBasedRecommender = ContentBasedRecommender
21
  class EndpointHandler:
22
  def __init__(self, path=""):
23
  model_path = os.path.join(path, "model.pkl")