1 files changed, 66 insertions, 0 deletions

diff --git a/reformat.py b/reformat.py
new file mode 100644
index 0000000..c21056c
--- /dev/null
+++ b/reformat.py
@@ -0,0 +1,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])