train_input_path: "data/kuairand-27k/s8000_c100/kuairand-27k-train-*.parquet"
eval_input_path: "data/kuairand-27k/s8000_c100/kuairand-27k-eval-*.parquet"
model_dir: "experiments/kuairand/dlrm_hstu"
train_config {
    sparse_optimizer {
        rowwise_adagrad_optimizer {
            lr: 0.001
        }
        constant_learning_rate {
        }
    }
    dense_optimizer {
        adam_optimizer {
            lr: 0.001
        }
        constant_learning_rate {
        }
    }
    num_epochs: 1
    save_checkpoints_epochs: 1
    mixed_precision: "BF16"
}
data_config {
    batch_size: 32
    dataset_type: ParquetDataset
    num_workers: 2
    fg_mode: FG_NONE
    odps_data_quota_name: ""
    label_fields: ["item_action_weight", "item_target_watchtime"]
}
feature_configs {
    sequence_id_feature {
        feature_name: "video_id"
        num_buckets: 10000000
        embedding_dim: 256
        embedding_name: "video_id_emb"
        data_type: "FP16"
    }
}
feature_configs {
    sequence_id_feature {
        feature_name: "item_video_id"
        num_buckets: 10000000
        embedding_dim: 256
        embedding_name: "video_id_emb"
        data_type: "FP16"
    }
}
feature_configs {
    sequence_raw_feature {
        feature_name: "action_timestamp"
    }
}
feature_configs {
    sequence_raw_feature {
        feature_name: "item_query_time"
    }
}
feature_configs {
    sequence_raw_feature {
        feature_name: "action_weight"
    }
}
feature_configs {
    sequence_raw_feature {
        feature_name: "watch_time"
    }
}
feature_configs {
    id_feature {
        feature_name: "user_id"
        num_buckets: 10000000
        embedding_dim: 256
        data_type: "FP16"
    }
}
feature_configs {
    id_feature {
        feature_name: "user_active_degree"
        num_buckets: 8
        embedding_dim: 256
        data_type: "FP16"
    }
}
feature_configs {
    id_feature {
        feature_name: "follow_user_num_range"
        num_buckets: 9
        embedding_dim: 256
        data_type: "FP16"
    }
}
feature_configs {
    id_feature {
        feature_name: "fans_user_num_range"
        num_buckets: 9
        embedding_dim: 256
        data_type: "FP16"
    }
}
feature_configs {
    id_feature {
        feature_name: "friend_user_num_range"
        num_buckets: 8
        embedding_dim: 256
        data_type: "FP16"
    }
}
feature_configs {
    id_feature {
        feature_name: "register_days_range"
        num_buckets: 8
        embedding_dim: 256
        data_type: "FP16"
    }
}
model_config {
    feature_groups {
        group_name: "contextual"
        feature_names: "user_id"
        feature_names: "user_active_degree"
        feature_names: "follow_user_num_range"
        feature_names: "fans_user_num_range"
        feature_names: "friend_user_num_range"
        feature_names: "register_days_range"
        group_type: DEEP
    }
    feature_groups {
        group_name: "uih"
        feature_names: "video_id"
        group_type: JAGGED_SEQUENCE
    }
    feature_groups {
        group_name: "candidate"
        feature_names: "item_video_id"
        group_type: JAGGED_SEQUENCE
    }
    feature_groups {
        group_name: "uih_action"
        feature_names: "action_weight"
        group_type: JAGGED_SEQUENCE
    }
    feature_groups {
        group_name: "uih_watchtime"
        feature_names: "watch_time"
        group_type: JAGGED_SEQUENCE
    }
    feature_groups {
        group_name: "uih_timestamp"
        feature_names: "action_timestamp"
        group_type: JAGGED_SEQUENCE
    }
    feature_groups {
        group_name: "candidate_timestamp"
        feature_names: "item_query_time"
        group_type: JAGGED_SEQUENCE
    }
    dlrm_hstu {
        hstu {
            stu {
                embedding_dim: 512
                num_heads: 4
                hidden_dim: 128
                attention_dim: 128
                output_dropout_ratio: 0.1
                use_group_norm: true
            }
            input_dropout_ratio: 0.2
            attn_num_layers: 3
            positional_encoder {
                num_position_buckets: 8192
                num_time_buckets: 2048
                use_time_encoding: true
            }
            input_preprocessor {
                contextual_preprocessor {
                    action_encoder {
                        simple_action_encoder {
                            action_embedding_dim: 8
                            action_weights: [1, 2, 4, 8, 16, 32, 64, 128]
                        }
                    }
                    action_mlp {
                        simple_mlp {
                            hidden_dim: 256
                        }
                    }
                    content_encoder {
                        slice_content_encoder {}
                    }
                    content_mlp {
                        simple_mlp {
                            hidden_dim: 256
                        }
                    }
                }
            }
            output_postprocessor {
                timestamp_layernorm_postprocessor {
                    time_duration_period_units: [3600, 86400]
                    time_duration_units_per_period: [24, 7]
                }
            }
        }
        fusion_mtl_tower {
            mlp {
                hidden_units: 512
                activation: "nn.SiLU"
                use_ln: true
            }
            task_configs {
                task_name: "is_click"
                label_name: "item_action_weight"
                task_bitmask: 1
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
            task_configs {
                task_name: "is_like"
                label_name: "item_action_weight"
                task_bitmask: 2
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
            task_configs {
                task_name: "is_follow"
                label_name: "item_action_weight"
                task_bitmask: 4
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
            task_configs {
                task_name: "is_comment"
                label_name: "item_action_weight"
                task_bitmask: 8
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
            task_configs {
                task_name: "is_forward"
                label_name: "item_action_weight"
                task_bitmask: 16
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
            task_configs {
                task_name: "is_hate"
                label_name: "item_action_weight"
                task_bitmask: 32
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
            task_configs {
                task_name: "long_view"
                label_name: "item_action_weight"
                task_bitmask: 64
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
            task_configs {
                task_name: "is_profile_enter"
                label_name: "item_action_weight"
                task_bitmask: 128
                losses {
                    binary_cross_entropy {}
                }
                metrics {
                    auc {}
                }
            }
        }
        max_seq_len: 8000
    }
    kernel: TRITON
}