main.py

import lancedb

import numpy as np
import pandas as pd

global data
data = []
global table
table = None


def get_recommendations(title):
    pd_data = pd.DataFrame(data)
    # Table Search
    result = (
        table.search(pd_data[pd_data["title"] == title]["vector"].values[0])
        .limit(5)
        .to_df()
    )

    # Get IMDB links
    links = pd.read_csv(
        "./ml-latest-small/links.csv",
        header=0,
        names=["movie id", "imdb id", "tmdb id"],
        converters={"imdb id": str},
    )

    ret = result["title"].values.tolist()
    # Loop to add links
    for i in range(len(ret)):
        link = links[links["movie id"] == result["id"].values[i]]["imdb id"].values[0]
        link = "https://www.imdb.com/title/tt" + link
        ret[i] = [ret[i], link]
    return ret


if __name__ == "__main__":
    # Load and prepare data
    ratings = pd.read_csv(
        "./ml-latest-small/ratings.csv",
        header=None,
        names=["user id", "movie id", "rating", "timestamp"],
    )
    ratings = ratings.drop(columns=["timestamp"])
    ratings = ratings.drop(0)
    ratings["rating"] = ratings["rating"].values.astype(np.float32)
    ratings["user id"] = ratings["user id"].values.astype(np.int32)
    ratings["movie id"] = ratings["movie id"].values.astype(np.int32)

    reviewmatrix = ratings.pivot(
        index="user id", columns="movie id", values="rating"
    ).fillna(0)

    # SVD
    matrix = reviewmatrix.values
    u, s, vh = np.linalg.svd(matrix, full_matrices=False)

    vectors = np.rot90(np.fliplr(vh))
    print(vectors.shape)

    # Metadata
    movies = pd.read_csv(
        "./ml-latest-small/movies.csv", header=0, names=["movie id", "title", "genres"]
    )
    movies = movies[movies["movie id"].isin(reviewmatrix.columns)]

    data = []
    for i in range(len(movies)):
        data.append(
            {
                "id": movies.iloc[i]["movie id"],
                "title": movies.iloc[i]["title"],
                "vector": vectors[i],
                "genre": movies.iloc[i]["genres"],
            }
        )
    print(pd.DataFrame(data))

    # Connect to LanceDB

    db = lancedb.connect("./data/test-db")
    try:
        table = db.create_table("movie_set", data=data)
    except:
        table = db.open_table("movie_set")

    print(get_recommendations("Moana (2016)"))
    print(get_recommendations("Rogue One: A Star Wars Story (2016)"))