wehub-resource-sync
104 lines
3.1 KiB
Python
104 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert training data JSONL from train_contrastors.py format to ms-swift embedding format.
|
|
|
|
Input format (train_contrastors.py):
|
|
{"query": "...", "chunk_path": "/path/to/pos.png"}
|
|
{"query": "...", "chunk_path": "/path/to/pos.png", "neg_chunk_paths": ["/path/to/neg1.png", ...]}
|
|
|
|
Output format (ms-swift embedding):
|
|
{
|
|
"messages": [
|
|
{"role": "system", "content": "Retrieve images or text relevant to the user's query."},
|
|
{"role": "user", "content": "<query>"}
|
|
],
|
|
"positive_messages": [[
|
|
{"role": "system", "content": "Represent the user's input."},
|
|
{"role": "user", "content": "<image>"}
|
|
]],
|
|
"positive_images": [["/path/to/pos.png"]],
|
|
"negative_messages": [
|
|
[{"role": "system", "content": "Represent the user's input."}, {"role": "user", "content": "<image>"}],
|
|
...
|
|
],
|
|
"negative_images": [["/path/to/neg1.png"], ...]
|
|
}
|
|
|
|
Usage:
|
|
uv run python convert_data_for_swift.py \
|
|
--input data/train_hn.jsonl --output data/train_hn_swift.jsonl
|
|
uv run python convert_data_for_swift.py \
|
|
--input data/eval.jsonl --output data/eval_swift.jsonl
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
|
|
# Match train_contrastors.py instructions exactly
|
|
QUERY_INSTRUCTION = "Retrieve images or text relevant to the user's query."
|
|
DOC_INSTRUCTION = "Represent the user's input."
|
|
|
|
|
|
def convert_line(item):
|
|
"""Convert one data item from contrastors format to swift format."""
|
|
query = item["query"]
|
|
pos_path = item["chunk_path"]
|
|
|
|
if not os.path.exists(pos_path):
|
|
return None
|
|
|
|
doc_messages = [
|
|
{"role": "system", "content": DOC_INSTRUCTION},
|
|
{"role": "user", "content": "<image>"},
|
|
]
|
|
|
|
out = {
|
|
"messages": [
|
|
{"role": "system", "content": QUERY_INSTRUCTION},
|
|
{"role": "user", "content": query},
|
|
],
|
|
"positive_messages": [doc_messages],
|
|
"positive_images": [[pos_path]],
|
|
}
|
|
|
|
neg_paths = item.get("neg_chunk_paths", [])
|
|
if neg_paths:
|
|
neg_messages = []
|
|
neg_images = []
|
|
for np_ in neg_paths:
|
|
if np_ and os.path.exists(np_):
|
|
neg_messages.append(doc_messages)
|
|
neg_images.append([np_])
|
|
if neg_messages:
|
|
out["negative_messages"] = neg_messages
|
|
out["negative_images"] = neg_images
|
|
|
|
return out
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--input", required=True, help="Input JSONL (contrastors format)"
|
|
)
|
|
parser.add_argument("--output", required=True, help="Output JSONL (swift format)")
|
|
args = parser.parse_args()
|
|
|
|
converted = 0
|
|
skipped = 0
|
|
with open(args.input) as fin, open(args.output, "w") as fout:
|
|
for line in fin:
|
|
item = json.loads(line.strip())
|
|
out = convert_line(item)
|
|
if out is None:
|
|
skipped += 1
|
|
continue
|
|
fout.write(json.dumps(out, ensure_ascii=False) + "\n")
|
|
converted += 1
|
|
|
|
print(f"Converted {converted} samples, skipped {skipped} → {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|