()
| 18 | from trl import SFTTrainer |
| 19 | |
| 20 | def get_data(): |
| 21 | data_folder = "data/agent_trajectories/nav" |
| 22 | all_json_files = glob.glob(os.path.join(data_folder, "**/*.json"), recursive=True) |
| 23 | system = [] |
| 24 | chat = [] |
| 25 | responses = [] |
| 26 | for file in all_json_files: |
| 27 | with open(file, "r") as f: |
| 28 | data = json.load(f) |
| 29 | system.append(data["system_prompt"]) |
| 30 | chat.append(data["human_message"]) |
| 31 | responses.append(data["system_response"]) |
| 32 | |
| 33 | def create_text_row(system_instruction, input, output): |
| 34 | text_row = f"""<s>[INST] <</SYS>>\\nSystem: {system_instruction}\\n<</SYS>>\\n\\nHuman: {input}[/INST] \nAssistant: \n{output}</s>""" |
| 35 | return text_row |
| 36 | |
| 37 | with open("data/hyperagent_traces.jsonl", "w") as f: |
| 38 | for prompt, input, output in zip(system, chat, responses): |
| 39 | object = { |
| 40 | "text": create_text_row(prompt, input, output) |
| 41 | } |
| 42 | json.dump(object, f) |
| 43 | f.write("\n") |
| 44 | |
| 45 | train_dataset = Dataset.from_dict({ |
| 46 | "text": [create_text_row(system_instruction, input, output) for system_instruction, input, output in zip(system, chat, responses)] |
| 47 | }) |
| 48 | return train_dataset |
| 49 | |
| 50 | def main(): |
| 51 | base_model_id = "" |
no test coverage detected