def extract_features(feature_vector):
"""Maps a feature vector to whether each feature is continuous or discrete."""
return [
DiscreteFeature(feature_vector[0]),
DiscreteFeature(feature_vector[1]),
ContinuousFeature(feature_vector[2])
]
define_flags()
FLAGS = flags.FLAGS
###############################################################################
# Load data
###############################################################################
def binarization(user_rating: float) -> float:
return 0.0 if user_rating <= 3.0 else 1.0
features_dict = build_features_dict([
StringFeature('movie_id'),
StringFeature('user_id'),
ContinuousFeature('label', 'user_rating', binarization),
IntegerFeature('user_gender', transform_fn=int),
StringFeature('user_zip_code'),
StringFeature('user_occupation_text'),
IntegerFeature('bucketized_user_age', transform_fn=int)
])
def load_data(features_dict):
dataset = f'movielens/{FLAGS.dataset}-ratings'
ratings = tfds.load(dataset, split='train', data_dir=FLAGS.data_dir)
# Prepare for binarization
ratings.filter(lambda x: x['user_rating'] != 3.0)
ratings = prepare_dataset(ratings, features_dict)
def extract_features(feature_vector):
"""Maps a feature vector to whether each feature is continuous or discrete."""
return [
DiscreteFeature(feature_vector[0]),
ContinuousFeature(feature_vector[1]),
ContinuousFeature(feature_vector[2]),
ContinuousFeature(feature_vector[3]),
ContinuousFeature(feature_vector[4]),
ContinuousFeature(feature_vector[5]),
ContinuousFeature(feature_vector[6]),
ContinuousFeature(feature_vector[7]),
ContinuousFeature(feature_vector[8]),
ContinuousFeature(feature_vector[9]),
ContinuousFeature(feature_vector[10]),
ContinuousFeature(feature_vector[11]),
ContinuousFeature(feature_vector[12]),
ContinuousFeature(feature_vector[13]),
ContinuousFeature(feature_vector[14]),
ContinuousFeature(feature_vector[15]),
ContinuousFeature(feature_vector[16]),
ContinuousFeature(feature_vector[17]),
ContinuousFeature(feature_vector[18]),
ContinuousFeature(feature_vector[19]),
ContinuousFeature(feature_vector[20]),
ContinuousFeature(feature_vector[21]),
ContinuousFeature(feature_vector[22]),
ContinuousFeature(feature_vector[23]),
ContinuousFeature(feature_vector[24]),
ContinuousFeature(feature_vector[25]),
ContinuousFeature(feature_vector[26]),
ContinuousFeature(feature_vector[27]),
ContinuousFeature(feature_vector[28]),
ContinuousFeature(feature_vector[29])
]
def extract_features(feature_vector):
return [
DiscreteFeature(feature_vector[0]),
DiscreteFeature(feature_vector[1]),
ContinuousFeature(feature_vector[2])
]