| 967 | |
| 968 | |
| 969 | class RLSid2TitleDataset(JSONBaseDataset): |
| 970 | def __init__(self, item_file, index_file, sample=-1, seed=0, category="", dedup=False): |
| 971 | """ |
| 972 | RL-specific dataset for sid2title tasks. |
| 973 | Returns prompt-completion pairs for RL training where input is semantic ID and output is item title. |
| 974 | |
| 975 | Args: |
| 976 | item_file: Path to .item.json file with item features |
| 977 | index_file: Path to .index.json file with item indices |
| 978 | sample: Number of samples to use (-1 for all) |
| 979 | seed: Random seed |
| 980 | category: Category name for prompts |
| 981 | dedup: Whether to filter duplicate items |
| 982 | """ |
| 983 | super().__init__(item_file, index_file, tokenizer=None, max_len=1024, test=False, category=category, dedup=dedup, seed=seed) |
| 984 | |
| 985 | self.prompt2history = {} |
| 986 | self.history2target = {} |
| 987 | |
| 988 | # Build sid2title mapping |
| 989 | self.sid2title = {} |
| 990 | |
| 991 | for item_id, sids in self.indices.items(): |
| 992 | if item_id in self.item_feat: |
| 993 | title = self.item_feat[item_id]['title'] |
| 994 | # Concatenate all three semantic IDs as the key |
| 995 | if len(sids) >= 3: |
| 996 | combined_sid = sids[0] + sids[1] + sids[2] |
| 997 | self.sid2title[combined_sid] = title |
| 998 | |
| 999 | # Create data samples |
| 1000 | self.data = [] |
| 1001 | |
| 1002 | # Create sid2title samples |
| 1003 | for sid, title in self.sid2title.items(): |
| 1004 | self.data.append({ |
| 1005 | 'task': 'sid2title', |
| 1006 | 'input': sid, |
| 1007 | 'output': title |
| 1008 | }) |
| 1009 | |
| 1010 | if sample > 0 and sample < len(self.data): |
| 1011 | self.data = random.sample(self.data, sample) |
| 1012 | |
| 1013 | self.get_inputs() |
| 1014 | |
| 1015 | def generate_prompt(self, data_point): |
| 1016 | prompt = f'What is the title of item "{data_point["input"]}"?' |
| 1017 | response = data_point['output'] |
| 1018 | |
| 1019 | return f"""### User Input: |
| 1020 | {prompt} |
| 1021 | |
| 1022 | ### Response:\n""" |
| 1023 | |
| 1024 | def pre(self, idx): |
| 1025 | data_point = self.data[idx] |
| 1026 | prompt = self.generate_prompt(data_point) |
no outgoing calls