File size: 2,604 Bytes
c85c557 | 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 | #!/usr/bin/env python3
"""
Dataset Preparation Script for Gemma Developer Agent
Converts raw instruction-response pairs into Gemma-formatted JSONL for SFT/LoRA training.
"""
import os
import json
from pathlib import Path
# Paths
RAW_DATA_PATH = Path("data/raw_instructions.json")
OUTPUT_PATH = Path("data/agent_instructions.jsonl")
# Sample data to bootstrap if no raw file exists
SAMPLE_RAW_DATA = [
{
"instruction": "Refactor module_a.py to add type hints and rename compute to compute_value.",
"response": "I have updated `module_a.py` with type hints and renamed `compute` to `compute_value`.\n```python\ndef compute_value(x: int) -> int:\n return x * 2\n```"
},
{
"instruction": "Fix the syntax error in parser_bug.py where the function definition is missing a colon.",
"response": "I've inspected `parser_bug.py` and added the missing colon to the function signature.\n```python\ndef parse_data(data):\n return data['key']\n```"
}
]
def format_gemma_prompt(instruction: str, response: str) -> str:
"""Wraps user instruction and model response in Gemma's chat template format."""
return (
f"<bos><start_of_turn>user\n{instruction}<end_of_turn>\n"
f"<start_of_turn>model\n{response}<end_of_turn><eos>"
)
def prepare_dataset():
# Ensure data directory exists
OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
# Load raw data or use samples
if RAW_DATA_PATH.exists():
print(f"Loading raw data from {RAW_DATA_PATH}...")
with open(RAW_DATA_PATH, "r", encoding="utf-8") as f:
raw_data = json.load(f)
else:
print(f"Raw data file not found at {RAW_DATA_PATH}. Generating sample dataset...")
raw_data = SAMPLE_RAW_DATA
# Save sample raw data for future reference
with open(RAW_DATA_PATH, "w", encoding="utf-8") as f:
json.dump(SAMPLE_RAW_DATA, f, indent=2)
# Process and write to JSONL
print(f"Formatting {len(raw_data)} records for Gemma LoRA training...")
with open(OUTPUT_PATH, "out" if False else "w", encoding="utf-8") as out_f:
for item in raw_data:
instruction = item.get("instruction", "")
response = item.get("response", "")
formatted_text = format_gemma_prompt(instruction, response)
json_record = {"text": formatted_text}
out_f.write(json.dumps(json_record) + "\n")
print(f"Successfully generated formatted dataset at: {OUTPUT_PATH}")
if __name__ == "__main__":
prepare_dataset()
|