1# /// script
2# dependencies = [
3#   "openai>=1.0",
4# ]
5# ///
6"""
7convert_dataset.py — Convert between SFT, DPO, and RFT dataset formats.
8 
9Usage:
10  # Parquet/CSV to SFT JSONL
11  python convert_dataset.py --input data.parquet --output train.jsonl --format sft \
12      --user-column prompt --assistant-column response --system-prompt "You are helpful."
13 
14  # SFT JSONL to DPO (generates rejected via base model)
15  python convert_dataset.py --input train.jsonl --output dpo.jsonl --format dpo \
16      --base-model gpt-4.1-mini --endpoint $ENDPOINT --api-key $KEY
17 
18  # SFT JSONL to RFT JSONL (passthrough — same format, different intent)
19  python convert_dataset.py --input train.jsonl --output rft.jsonl --format rft
20 
21  # DPO JSONL to SFT (extract chosen responses)
22  python convert_dataset.py --input dpo.jsonl --output sft.jsonl --format sft-from-dpo
23"""
24 
25import json
26import os
27import sys
28 
29try:
30    sys.stdout.reconfigure(encoding="utf-8")
31    sys.stderr.reconfigure(encoding="utf-8")
32except (AttributeError, OSError):
33    pass  # Stream not reconfigurable (older Python or non-tty); default encoding is fine
34import time
35sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
36from common import HelpOnErrorParser, get_clients
37 
38 
39def parquet_to_sft(input_path, output_path, user_col, assistant_col, system_prompt=None):
40    """Convert a parquet or CSV file to SFT JSONL."""
41    try:
42        import pandas as pd
43    except ImportError:
44        print("Error: pandas required. Install with: pip install pandas pyarrow")
45        sys.exit(1)
46 
47    if input_path.endswith(".parquet"):
48        df = pd.read_parquet(input_path)
49    elif input_path.endswith(".csv"):
50        df = pd.read_csv(input_path)
51    elif input_path.endswith(".json"):
52        df = pd.read_json(input_path)
53    else:
54        print(f"Unsupported format: {input_path}. Use .parquet, .csv, or .json")
55        sys.exit(1)
56 
57    if user_col not in df.columns or assistant_col not in df.columns:
58        print(f"Error: Columns '{user_col}' and/or '{assistant_col}' not found.")
59        print(f"Available columns: {list(df.columns)}")
60        sys.exit(1)
61 
62    count = 0
63    with open(output_path, "w", encoding="utf-8") as f:
64        for _, row in df.iterrows():
65            user_content = str(row[user_col]).strip()
66            asst_content = str(row[assistant_col]).strip()
67            if not user_content or not asst_content:
68                continue
69 
70            messages = []
71            if system_prompt:
72                messages.append({"role": "system", "content": system_prompt})
73            messages.append({"role": "user", "content": user_content})
74            messages.append({"role": "assistant", "content": asst_content})
75 
76            f.write(json.dumps({"messages": messages}, ensure_ascii=False) + "\n")
77            count += 1
78 
79    print(f"Converted {count} examples to SFT JSONL → {output_path}")
80 
81 
82def sft_to_dpo(input_path, output_path, client, base_model):
83    """Convert SFT to DPO by generating non-preferred responses from a base model.
84 
85    DPO format uses: input (system+user messages), preferred_output, non_preferred_output.
86    """
87    with open(input_path, encoding="utf-8") as inf:
88        examples = []
89        for ln, raw in enumerate(inf, 1):
90            if not raw.strip():
91                continue
92            try:
93                examples.append(json.loads(raw))
94            except json.JSONDecodeError as e:
95                print(f"  ⚠️ Skipping malformed JSON on line {ln}: {e}")
96    count = 0
97 
98    with open(output_path, "w", encoding="utf-8") as f:
99        for i, ex in enumerate(examples):
100            msgs = ex["messages"]
101            system_msgs = [m for m in msgs if m["role"] == "system"]
102            user_msg = next((m for m in msgs if m["role"] == "user"), None)
103            asst_msg = next((m for m in msgs if m["role"] == "assistant"), None)
104            if not user_msg or not asst_msg:
105                continue
106 
107            # Generate a non-preferred response from the base model
108            try:
109                gen_msgs = system_msgs + [user_msg]
110                resp = client.chat.completions.create(
111                    model=base_model,
112                    messages=gen_msgs,
113                    temperature=1.0,  # High temp for diversity
114                    max_completion_tokens=2048,
115                )
116                rejected_content = resp.choices[0].message.content
117            except Exception as e:
118                print(f"  Skipping example {i}: {e}")
119                continue
120 
121            if not rejected_content:
122                # None or empty — content filter, finish=length with no text, etc.
123                # Skip rather than emit a DPO entry with null content (trainer rejects).
124                print(f"  Skipping example {i}: base model returned no content")
125                continue
126 
127            # Build DPO entry with correct format
128            input_messages = system_msgs + [user_msg]
129            dpo_entry = {
130                "input": {"messages": input_messages},
131                "preferred_output": [asst_msg],
132                "non_preferred_output": [{"role": "assistant", "content": rejected_content}],
133            }
134            f.write(json.dumps(dpo_entry, ensure_ascii=False) + "\n")
135            count += 1
136 
137            if (i + 1) % 50 == 0:
138                print(f"  Processed {i+1}/{len(examples)}")
139                time.sleep(1)
140 
141    print(f"Converted {count} examples to DPO JSONL → {output_path}")
142 
143 
144def sft_to_rft(input_path, output_path):
145    """Convert SFT to RFT format.
146 
147    Strips assistant messages (RFT last message must be user) and adds a
148    placeholder grader field. The user must populate grader reference fields
149    (e.g., expected_answer) before training.
150    """
151    count = 0
152    skipped = 0
153    with open(output_path, "w", encoding="utf-8") as out:
154        with open(input_path, encoding="utf-8") as inf:
155            for ln, line in enumerate(inf, 1):
156                if not line.strip():
157                    continue
158                try:
159                    ex = json.loads(line)
160                except json.JSONDecodeError as e:
161                    print(f"  ⚠️ Skipping malformed JSON on line {ln}: {e}")
162                    skipped += 1
163                    continue
164                msgs = ex.get("messages", [])
165                # Keep only system + user messages; RFT last message must be user
166                rft_msgs = [m for m in msgs if m["role"] in ("system", "user")]
167                if not rft_msgs or rft_msgs[-1]["role"] != "user":
168                    skipped += 1
169                    continue
170                # Extract assistant content as a reference answer placeholder
171                asst_msgs = [m for m in msgs if m["role"] == "assistant"]
172                expected = asst_msgs[-1]["content"] if asst_msgs else ""
173                rft_entry = {"messages": rft_msgs, "expected_answer": expected}
174                out.write(json.dumps(rft_entry, ensure_ascii=False) + "\n")
175                count += 1
176    print(f"Converted {count} examples to RFT JSONL → {output_path}")
177    if skipped:
178        print(f"  Skipped {skipped} examples (no user message)")
179    print("Note: Review 'expected_answer' fields and update your grader to use item.expected_answer.")
180 
181 
182def dpo_to_sft(input_path, output_path, system_prompt=None):
183    """Extract chosen responses from DPO format to SFT format."""
184    count = 0
185    with open(output_path, "w", encoding="utf-8") as f:
186        with open(input_path, encoding="utf-8") as inf:
187            for ln, line in enumerate(inf, 1):
188                if not line.strip():
189                    continue
190                try:
191                    ex = json.loads(line)
192                except json.JSONDecodeError as e:
193                    print(f"  ⚠️ Skipping malformed JSON on line {ln}: {e}")
194                    continue
195                input_messages = ex["input"]["messages"]
196                chosen_messages = ex["preferred_output"]
197 
198                messages = []
199                if system_prompt:
200                    messages.append({"role": "system", "content": system_prompt})
201                    messages.extend(m for m in input_messages if m["role"] != "system")
202                else:
203                    messages.extend(input_messages)
204                messages.extend(chosen_messages)
205                f.write(json.dumps({"messages": messages}, ensure_ascii=False) + "\n")
206                count += 1
207    print(f"Extracted {count} chosen examples to SFT JSONL → {output_path}")
208 
209 
210def main():
211    parser = HelpOnErrorParser(description="Convert between fine-tuning dataset formats")
212    parser.add_argument("--input", required=True, help="Input file path")
213    parser.add_argument("--output", required=True, help="Output file path")
214    parser.add_argument("--format", required=True,
215                        choices=["sft", "dpo", "rft", "sft-from-dpo"],
216                        help="Target format")
217 
218    # SFT from raw data
219    parser.add_argument("--user-column", default="prompt", help="Column name for user input")
220    parser.add_argument("--assistant-column", default="response", help="Column name for assistant output")
221    parser.add_argument("--system-prompt", default=None, help="System prompt to prepend")
222 
223    # DPO generation (needs API connection)
224    parser.add_argument("--base-url", default=os.environ.get("OPENAI_BASE_URL"),
225                        help="Project /v1/ URL (preferred)")
226    parser.add_argument("--endpoint", default=os.environ.get("AZURE_OPENAI_ENDPOINT"),
227                        help="Azure OpenAI endpoint (fallback)")
228    parser.add_argument("--project-endpoint", default=os.environ.get("AZURE_AI_PROJECT_ENDPOINT"),
229                        help="Azure AI project endpoint (Foundry SDK)")
230    parser.add_argument("--api-key", default=os.environ.get("AZURE_OPENAI_API_KEY"))
231    parser.add_argument("--base-model", default="gpt-4.1-mini", help="Base model for generating rejections")
232 
233    args = parser.parse_args()
234 
235    if args.format == "sft":
236        if args.input.endswith(".jsonl"):
237            print("Input is already JSONL — assuming SFT format. Nothing to convert.")
238            if args.input != args.output:
239                import shutil
240                shutil.copy2(args.input, args.output)
241        else:
242            parquet_to_sft(args.input, args.output, args.user_column,
243                           args.assistant_column, args.system_prompt)
244 
245    elif args.format == "dpo":
246        client, method = get_clients(
247            base_url=args.base_url, azure_endpoint=args.endpoint,
248            project_endpoint=args.project_endpoint, api_key=args.api_key
249        )
250        sft_to_dpo(args.input, args.output, client, args.base_model)
251 
252    elif args.format == "rft":
253        sft_to_rft(args.input, args.output)
254 
255    elif args.format == "sft-from-dpo":
256        dpo_to_sft(args.input, args.output, args.system_prompt)
257 
258 
259if __name__ == "__main__":
260    main()
261