(
# model/data params
model_path: str = "",
seed: int = 42,
train_file: str = "",
eval_file: str = "",
info_file: str = "",
category: str = "",
# wandb params
wandb_project: str = "",
wandb_run_name: str = "",
# training hyperparams
output_dir: str = "",
train_batch_size: int = 32,
eval_batch_size: int = 32,
gradient_accumulation_steps: int = 1,
temperature: float = 1.0,
add_gt: bool = False,
eval_step: float = 0.199,
num_generations: int = 16,
num_train_epochs: int = 1,
learning_rate: float = 1e-6,
beta: float = 0.04,
beam_search: bool = False,
test_during_training: bool = True,
dynamic_sampling: bool = False,
mask_all_zero: bool = False,
sync_ref_model: bool = False,
test_beam: int = 20,
reward_type: str = "rule",
sample_train: bool = False,
ada_path: str = "",
cf_path: str = "",
sid_index_path: str = "",
item_meta_path: str = "",
dapo: bool = False,
gspo: bool = False,
)
| 28 | torch.backends.cudnn.benchmark = False |
| 29 | |
| 30 | def train( |
| 31 | # model/data params |
| 32 | model_path: str = "", |
| 33 | seed: int = 42, |
| 34 | train_file: str = "", |
| 35 | eval_file: str = "", |
| 36 | info_file: str = "", |
| 37 | category: str = "", |
| 38 | |
| 39 | # wandb params |
| 40 | wandb_project: str = "", |
| 41 | wandb_run_name: str = "", |
| 42 | |
| 43 | # training hyperparams |
| 44 | output_dir: str = "", |
| 45 | train_batch_size: int = 32, |
| 46 | eval_batch_size: int = 32, |
| 47 | gradient_accumulation_steps: int = 1, |
| 48 | temperature: float = 1.0, |
| 49 | add_gt: bool = False, |
| 50 | eval_step: float = 0.199, |
| 51 | num_generations: int = 16, |
| 52 | num_train_epochs: int = 1, |
| 53 | learning_rate: float = 1e-6, |
| 54 | beta: float = 0.04, |
| 55 | beam_search: bool = False, |
| 56 | test_during_training: bool = True, |
| 57 | dynamic_sampling: bool = False, |
| 58 | mask_all_zero: bool = False, |
| 59 | sync_ref_model: bool = False, |
| 60 | test_beam: int = 20, |
| 61 | reward_type: str = "rule", |
| 62 | sample_train: bool = False, |
| 63 | ada_path: str = "", |
| 64 | cf_path: str = "", |
| 65 | sid_index_path: str = "", |
| 66 | item_meta_path: str = "", |
| 67 | dapo: bool = False, |
| 68 | gspo: bool = False, |
| 69 | ): |
| 70 | torch.backends.cuda.enable_flash_sdp(False) |
| 71 | torch.backends.cuda.enable_mem_efficient_sdp(False) |
| 72 | set_seed(seed) |
| 73 | |
| 74 | category_dict = {"Industrial_and_Scientific": "industrial and scientific items", "Office_Products": "office products", "Toys_and_Games": "toys and games", "Sports": "sports and outdoors", "Books": "books"} |
| 75 | print(category) |
| 76 | |
| 77 | |
| 78 | with open(info_file, 'r') as f: |
| 79 | info = f.readlines() |
| 80 | # Extract semantic_id (first column) from the format: semantic_id \t item_title \t item_id |
| 81 | item_name = [_.split('\t')[0].strip() for _ in info] |
| 82 | item2id = {name: i for i, name in enumerate(item_name)} |
| 83 | |
| 84 | sample = -1 |
| 85 | train_datasets = [] |
| 86 | # train_data = D3Dataset(train_file, category=category_dict[category], sample=sample) |
| 87 | # train_datasets.append(train_data) |
nothing calls this directly
no test coverage detected