50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
import json
|
|
import os
|
|
import random
|
|
from typing import AnyStr
|
|
|
|
# parsing OA data files with oasst_data helpers
|
|
from oasst_data import read_message_trees, ExportMessageNode
|
|
|
|
messages: list[ExportMessageNode] = []
|
|
|
|
this_dir = os.path.dirname(os.path.abspath(__file__))
|
|
input_file_path = this_dir + '/2023-11-05_oasst2_all.trees.jsonl.gz'
|
|
|
|
role_dict = {'prompter': 'user', 'assistant': 'assistant'}
|
|
user_names = ['Adam', 'Alice', 'Anne', 'Bob', 'Charlie', 'Cody', 'Corinna', 'Cynthia', 'Fred', 'Grace', 'Jane', 'Paul',
|
|
'Rachel', 'Ramesh']
|
|
conversations = []
|
|
|
|
|
|
def visit(node: ExportMessageNode, parents: [ExportMessageNode], user: AnyStr):
|
|
new_parents = parents + [node]
|
|
if not node.replies: # end of conversation
|
|
conversations.append({'messages': [{'role': role_dict[p.role],
|
|
'name': user if role_dict[p.role] != 'assistant' else 'assistant',
|
|
'content': p.text
|
|
} for p in new_parents]})
|
|
else:
|
|
for reply in node.replies:
|
|
visit(reply, new_parents, user)
|
|
|
|
|
|
for tree in read_message_trees(input_file_path):
|
|
if tree.prompt.lang not in ['en']: # filtering by language tag (optional)
|
|
continue
|
|
|
|
visit(tree.prompt, [], user_names[random.randint(0, len(user_names) - 1)])
|
|
|
|
print(conversations[0])
|
|
|
|
|
|
def mkdir(path):
|
|
if not os.path.isdir(path):
|
|
os.mkdir(path)
|
|
|
|
|
|
mkdir(this_dir + '/../../data')
|
|
mkdir(this_dir + '/../../data/oasst')
|
|
with open(this_dir + '/../../data/oasst/oasst_all.jsonl', 'w') as f:
|
|
f.writelines([json.dumps(conv) + '\n' for conv in conversations])
|