(
base_model: str = "",
train_file: str = "",
info_file: str = "",
category: str = "",
test_data_path: str = "",
result_json_data: str = "",
batch_size: int = 4,
K: int = 0,
seed: int = 42,
length_penalty: float=0.0,
max_new_tokens: int = 256,
num_beams: int = 50,
)
| 36 | # torch.backends.cudnn.benchmark = False |
| 37 | |
| 38 | def main( |
| 39 | base_model: str = "", |
| 40 | train_file: str = "", |
| 41 | info_file: str = "", |
| 42 | category: str = "", |
| 43 | test_data_path: str = "", |
| 44 | result_json_data: str = "", |
| 45 | batch_size: int = 4, |
| 46 | K: int = 0, |
| 47 | seed: int = 42, |
| 48 | length_penalty: float=0.0, |
| 49 | max_new_tokens: int = 256, |
| 50 | num_beams: int = 50, |
| 51 | ): |
| 52 | random.seed(seed) |
| 53 | set_seed(seed) |
| 54 | os.environ["CUDA_VISIBLE_DEVICES"] = "0" |
| 55 | 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"} |
| 56 | category = category_dict[category] |
| 57 | print(category) |
| 58 | |
| 59 | model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.bfloat16, device_map="auto") |
| 60 | model.eval() |
| 61 | with open(info_file, 'r') as f: |
| 62 | info = f.readlines() |
| 63 | # Parse new format: semantic_id \t item_title \t item_id |
| 64 | semantic_ids = [line.split('\t')[0].strip() + "\n" for line in info] |
| 65 | item_titles = [line.split('\t')[1].strip() + "\n" for line in info if len(line.split('\t')) >= 2] |
| 66 | |
| 67 | # Format for tokenization |
| 68 | info_semantic = [f'''### Response:\n{_}''' for _ in semantic_ids] |
| 69 | info_titles = [f'''### Response:\n{_}''' for _ in item_titles] |
| 70 | |
| 71 | |
| 72 | tokenizer = AutoTokenizer.from_pretrained(base_model) |
| 73 | |
| 74 | # Create prefixID for semantic IDs (existing functionality) |
| 75 | if base_model.lower().find("llama") > -1: |
| 76 | prefixID = [tokenizer(_).input_ids[1:] for _ in info_semantic] |
| 77 | prefixTitleID = [tokenizer(_).input_ids[1:] for _ in info_titles] |
| 78 | else: |
| 79 | prefixID = [tokenizer(_).input_ids for _ in info_semantic] |
| 80 | prefixTitleID = [tokenizer(_).input_ids for _ in info_titles] |
| 81 | if base_model.lower().find("gpt2") > -1: |
| 82 | prefix_index = 4 |
| 83 | else: |
| 84 | prefix_index = 3 |
| 85 | |
| 86 | # Build hash_dict for semantic IDs (existing functionality) |
| 87 | hash_dict = dict() |
| 88 | # print(f"eos token: {tokenizer.eos_token_id}") |
| 89 | for index, ID in enumerate(prefixID): |
| 90 | ID.append(tokenizer.eos_token_id) |
| 91 | for i in range(prefix_index, len(ID)): |
| 92 | if i == prefix_index: |
| 93 | hash_number = get_hash(ID[:i]) |
| 94 | else: |
| 95 | hash_number = get_hash(ID[prefix_index:i]) |
nothing calls this directly
no test coverage detected