Spaces:
Sleeping
Sleeping
File size: 2,762 Bytes
cd5a7c4 |
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 |
import pandas as pd
from qdrant_client import QdrantClient, models
from collections import defaultdict
# Load the data
tags = pd.read_csv("./data/ml-latest-small/tags.csv")
movies = pd.read_csv("./data/ml-latest-small/movies.csv")
ratings = pd.read_csv("./data/ml-latest-small/ratings.csv")
# Initialize Qdrant client and create collections
def init_qdrant():
qdrant = QdrantClient(":memory:") # Use in-memory for simplicity
qdrant.create_collection(
"movielens", vectors_config={}, sparse_vectors_config={"ratings": models.SparseVectorParams()}
)
return qdrant
# Load data and upload to Qdrant
def load_data(qdrant):
ratings['normalized_rating'] = (ratings.rating - ratings.rating.mean(axis=0)) / ratings.rating.std()
user_sparse_vectors = defaultdict(lambda: {"values": [], "indices": []})
for row in ratings.itertuples():
user_sparse_vectors[row.userId]["values"].append(row.normalized_rating)
user_sparse_vectors[row.userId]["indices"].append(row.movieId)
def data_generator():
for user_id, vector in user_sparse_vectors.items():
yield models.PointStruct(
id=user_id, vector={"ratings": vector}, payload={}
)
qdrant.upload_points("movielens", data_generator())
# Function to input and normalize ratings
def input_ratings(user_ratings, ratings):
final_ratings = {}
mean_rating = ratings.rating.mean()
std_rating = ratings.rating.std()
for movie_id, user_rating in user_ratings.values():
normalized_input_rating = (user_rating - mean_rating) / std_rating
final_ratings[movie_id] = normalized_input_rating
return final_ratings
# Search and recommendation function
def recommend_movies(qdrant, movies, my_ratings):
def to_vector(ratings):
vector = models.SparseVector(values=[], indices=[])
for movieId, rating in ratings.items():
vector.values.append(rating)
vector.indices.append(movieId)
return vector
user_vector = to_vector(my_ratings)
results = qdrant.search(
"movielens",
query_vector=models.NamedSparseVector(name="ratings", vector=user_vector),
with_vectors=True,
limit=20,
)
movie_scores = defaultdict(lambda: 0)
for user in results:
user_scores = user.vector["ratings"]
for idx, rating in zip(user_scores.indices, user_scores.values):
if idx in my_ratings:
continue
movie_scores[idx] += rating
top_movies = sorted(movie_scores.items(), key=lambda x: x[1], reverse=True)
recommended_movies = [movies[movies.movieId == movieId].title.values[0] for movieId, score in top_movies[:5]]
return recommended_movies
|