TensorFlow推荐系统（一）¶

# Dependency install
!pip install -q tensorflow-recommenders
!pip install -q --upgrade tensorflow-datasets

import os
import pprint
import tempfile
from typing import Dict, Text
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
## import tensorflow recommenders API
import tensorflow_recommenders as tfrs

# Ratings data.
ratings = tfds.load("movielens/100k-ratings", split="train")
# Features of all the available movies.
movies = tfds.load("movielens/100k-movies", split="train")
## Note: MovieLens 没有事先准备好数据集，所有的数据都在train data中。
# 查看数据rating和movies
for x in ratings.take(2).as_numpy_iterator():
  pprint.pprint(x)
#{'movie_title': b"One Flew Over the Cuckoo's Nest (1975)", 'user_id': b'138'}
#{'movie_title': b'Strictly Ballroom (1992)', 'user_id': b'92'}

for x in movies.take(2).as_numpy_iterator():
  pprint.pprint(x)
#b'You So Crazy (1994)'
#b'Love Is All There Is (1996)'

ratings = ratings.map(lambda x: {
    "movie_title": x["movie_title"],
    "user_id": x["user_id"],
})
movies = movies.map(lambda x: x["movie_title"])

tf.random.set_seed(42)
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)

train = shuffled.take(80_000)
test = shuffled.skip(80_000).take(20_000)
# 为变量数值做mapping
movie_titles = movies.batch(1_000)
user_ids = ratings.batch(1_000_000).map(lambda x: x["user_id"])

unique_movie_titles = np.unique(np.concatenate(list(movie_titles)))
unique_user_ids = np.unique(np.concatenate(list(user_ids)))

unique_movie_titles[:10]

#array([b"'Til There Was You (1997)", b'1-900 (1994)',
       #b'101 Dalmatians (1996)', b'12 Angry Men (1957)', b'187 (1997)',
       #b'2 Days in the Valley (1996)',
       #b'20,000 Leagues Under the Sea (1954)',
       #b'2001: A Space Odyssey (1968)',
       #b'3 Ninjas: High Noon At Mega Mountain (1998)',
       #b'39 Steps, The (1935)'], dtype=object)

#设定embedding维度
embedding_dimension = 32
#设定用户模型
user_model = tf.keras.Sequential([
  tf.keras.layers.StringLookup(
      vocabulary=unique_user_ids, mask_token=None),
  # We add an additional embedding to account for unknown tokens.
  tf.keras.layers.Embedding(len(unique_user_ids) + 1, embedding_dimension)
])
#设定电影模型
movie_model = tf.keras.Sequential([
  tf.keras.layers.StringLookup(
      vocabulary=unique_movie_titles, mask_token=None),
  tf.keras.layers.Embedding(len(unique_movie_titles) + 1, embedding_dimension)
])

metrics = tfrs.metrics.FactorizedTopK(
  candidates=movies.batch(128).map(movie_model)
)

task = tfrs.tasks.Retrieval(
  metrics=metrics
)

class MovielensModel(tfrs.Model):

  def __init__(self, user_model, movie_model):
    super().__init__()
    self.movie_model: tf.keras.Model = movie_model
    self.user_model: tf.keras.Model = user_model
    self.task: tf.keras.layers.Layer = task

  def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:
    # We pick out the user features and pass them into the user model.
    user_embeddings = self.user_model(features["user_id"])
    # And pick out the movie features and pass them into the movie model,
    # getting embeddings back.
    positive_movie_embeddings = self.movie_model(features["movie_title"])

    # The task computes the loss and the metrics.
    return self.task(user_embeddings, positive_movie_embeddings)

model = MovielensModel(user_model, movie_model)
model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))
#shuffle data and get samples
cached_train = train.shuffle(100_000).batch(8192).cache()
cached_test = test.batch(4096).cache()

model.fit(cached_train, epochs=3)
#使用test data evaluate
model.evaluate(cached_test, return_dict=True)
{'factorized_top_k/top_1_categorical_accuracy': 0.00044999999227002263,
 'factorized_top_k/top_5_categorical_accuracy': 0.004100000020116568,
 'factorized_top_k/top_10_categorical_accuracy': 0.01145000010728836,
 'factorized_top_k/top_50_categorical_accuracy': 0.09040000289678574,
 'factorized_top_k/top_100_categorical_accuracy': 0.19300000369548798,
 'loss': 28535.75390625,
 'regularization_loss': 0,
 'total_loss': 28535.75390625}

# Create a model that takes in raw query features, and
index = tfrs.layers.factorized_top_k.BruteForce(model.user_model)
# recommends movies out of the entire movies dataset.
index.index_from_dataset(
  tf.data.Dataset.zip((movies.batch(100), movies.batch(100).map(model.movie_model)))
)
# Get recommendations.
_, titles = index(tf.constant(["9"]))
print(f"Recommendations for user 9: {titles[0, :3]}")