def gen_class_weights(df):
  """Generate class weights to even out the class distribution during training."""
  class_counts_df = df.select("tumor_score").groupBy("tumor_score").count()
  class_counts = {row["tumor_score"]:row["count"] for row in class_counts_df.collect()}
  max_count = max(class_counts.values())
  class_weights = {k-1:max_count/v for k,v in class_counts.items()}
  return class_weights