1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
|
import sys
import json
from transformers import AutoTokenizer
from tqdm import tqdm
# IMPORTANT: The tokenizer MUST match the student model you are fine-tuning.
STUDENT_MODEL_ID = "google/t5gemma-s-s-ul2-it"
def reformat_with_template(input_filepath: str, output_filepath: str):
"""
Reads a JSONL file with 'text' (prompt) and 'target' (completion) fields
and reformats it into a single 'text' field using the model's chat template.
"""
print(f"Loading tokenizer for {STUDENT_MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL_ID)
print(f"Reformatting {input_filepath}...")
with open(input_filepath, 'r', encoding='utf-8') as infile, \
open(output_filepath, 'w', encoding='utf-8') as outfile:
# Count lines for tqdm progress bar
num_lines = sum(1 for line in open(input_filepath, 'r', encoding='utf-8'))
infile.seek(0) # Reset file pointer
for line in tqdm(infile, total=num_lines, desc="Reformatting"):
original_entry = json.loads(line)
# The original prompt is in the 'text' field
prompt_text = original_entry.get("text", "")
# The verified completion is in the 'target' field
completion_text = original_entry.get("target", "")
# The prompt needs to end with the cue for the model to respond to
if "Timestamp:" not in prompt_text:
prompt_text += "\n\nTimestamp:\n"
# Create the message structure that the chat template expects
messages = [
{"role": "user", "content": prompt_text},
{"role": "assistant", "content": completion_text}
]
# Apply the template to create the new, single text field
# add_generation_prompt=False is important for training data
formatted_text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=False
)
new_entry = {
"id": original_entry.get("id"),
"text": formatted_text,
"target": completion_text # Keep target for reference if needed
}
outfile.write(json.dumps(new_entry) + "\n")
print(f"\nSuccessfully reformatted {num_lines} entries. Output saved to {output_filepath}.")
if __name__ == "__main__":
if len(sys.argv) < 3:
print("Usage: python reformat_with_template.py <input_dataset.jsonl> <output_dataset_templated.jsonl>")
sys.exit(1)
reformat_with_template(sys.argv[1], sys.argv[2])
|
