adds improved launch for reasoning gpt-oss configs and new config for medical reasoning

config/train_gpt_oss_custom.py

@@ -109,6 +109,9 @@ class GPTOSSEnhancedCustomConfig:
     # Field Mapping - Customize for your dataset format
     input_field: str = "prompt"  # Field containing the input/prompt
     target_field: str = "accepted_completion"  # Field containing the target/completion
+    # Optional global conversational context
+    system_message: Optional[str] = None
+    developer_message: Optional[str] = None
     
     # OpenHermes-FR specific fields
     filter_bad_entries: bool = True  # Filter entries marked as bad
@@ -127,7 +130,14 @@ class GPTOSSEnhancedCustomConfig:
     max_length: Optional[int] = None  # Maximum sequence length (None = use max_seq_length)
     
     # Custom Dataset Formats Support
-    dataset_format: str = "openhermes_fr"  # "openhermes_fr", "messages", "text", "custom"
+    dataset_format: str = "openhermes_fr"  # "openhermes_fr", "messages", "text", "custom", "medical_o1_sft", "preference"
+
+    # Medical o1 SFT (FreedomIntelligence/medical-o1-reasoning-SFT) mapping
+    question_field: str = "Question"
+    reasoning_field: str = "Complex_CoT"
+    response_field: str = "Response"
+    reason_prefix: str = "Reasoning: "
+    answer_prefix: str = "Final Answer: "
     
     # GPT-OSS Harmony Format Configuration
     use_harmony_format: bool = True  # Enable GPT-OSS harmony format
@@ -344,7 +354,7 @@ class GPTOSSEnhancedCustomConfig:
             raise ValueError("max_seq_length must be >= 1")
             
         # Validate dataset format
-        valid_formats = ["openhermes_fr", "messages", "text", "custom"]
+        valid_formats = ["openhermes_fr", "messages", "text", "custom", "medical_o1_sft", "preference"]
         if self.dataset_format not in valid_formats:
             raise ValueError(f"dataset_format must be one of {valid_formats}")
     
@@ -383,6 +393,12 @@ class GPTOSSEnhancedCustomConfig:
         print(f"   • Target Field: {self.target_field}")
         print(f"   • Filter Bad Entries: {self.filter_bad_entries}")
         print(f"   • Max Samples: {self.max_samples or 'All'}")
+        if self.system_message or self.developer_message:
+            print("   • Context messages set:")
+            if self.system_message:
+                print("     - system message: provided")
+            if self.developer_message:
+                print("     - developer message: provided")
         
         print(f"\n💾 Memory & Performance:")
         print(f"   • Mixed Precision: {'BF16' if self.bf16 else 'FP32'}")

config/train_gpt_oss_medical_o1_sft.py

@@ -0,0 +1,151 @@
+"""
+GPT-OSS Medical o1 SFT Training Configuration
+Dataset: FreedomIntelligence/medical-o1-reasoning-SFT
+Format: Question | Complex_CoT | Response → GPT-OSS Harmony text
+
+This configuration uses GPT-OSS Harmony formatting to combine the medical
+dataset's question, chain-of-thought (Complex_CoT), and final response into a
+single assistant turn, with optional system and developer messages.
+"""
+
+from config.train_gpt_oss_custom import GPTOSSEnhancedCustomConfig
+
+# Medical-o1 SFT configuration for GPT-OSS
+config = GPTOSSEnhancedCustomConfig(
+    # ============================================================================
+    # DATASET CONFIGURATION
+    # ============================================================================
+    dataset_name="FreedomIntelligence/medical-o1-reasoning-SFT",
+    dataset_config="en",               # Use English split by default (can be changed to en_mix/zh/zh_mix)
+    dataset_split="train",
+    dataset_format="medical_o1_sft",   # Enable medical formatter in training script
+
+    # Field mapping and prefixes
+    input_field="Question",            # used for length filtering pre-format
+    target_field="Response",           # used for length filtering pre-format
+    question_field="Question",
+    reasoning_field="Complex_CoT",
+    response_field="Response",
+    reason_prefix="Reasoning: ",
+    answer_prefix="Final Answer: ",
+
+    # GPT-OSS Harmony formatting
+    use_harmony_format=True,
+    use_chat_template=False,
+    system_message=(
+        "You are GPT-Tonic, a large language model trained by TonicAI."
+    ),
+    developer_message=(
+        "You are an intelligent assistant that can answer customer service queries"
+    ),
+    chat_template_kwargs={
+        "add_generation_prompt": True,
+        "tokenize": False,
+        "reasoning_effort": "low",
+        "model_identity": "You are GPT-Tonic, a large language model trained by TonicAI.",
+        "builtin_tools": [],
+    },
+
+    # Filtering & sampling
+    filter_bad_entries=False,
+    max_samples=None,
+    min_length=10,
+    max_length=2048,
+
+    # ============================================================================
+    # TRAINING HYPERPARAMETERS
+    # ============================================================================
+    num_train_epochs=1.0,
+    batch_size=2,
+    gradient_accumulation_steps=8,
+    learning_rate=2e-4,
+    min_lr=2e-5,
+    weight_decay=0.01,
+    warmup_ratio=0.03,
+    max_grad_norm=1.0,
+
+    # Sequence length
+    max_seq_length=2048,
+
+    # ============================================================================
+    # MIXED PRECISION / PERFORMANCE
+    # ============================================================================
+    fp16=False,
+    bf16=True,
+    tf32=True,
+
+    dataloader_num_workers=4,
+    dataloader_pin_memory=True,
+    dataloader_prefetch_factor=2,
+    dataset_num_proc=4,
+    group_by_length=True,
+    remove_unused_columns=True,
+
+    # ============================================================================
+    # LORA & QUANTIZATION
+    # ============================================================================
+    use_lora=True,
+    lora_config={
+        "r": 8,
+        "lora_alpha": 16,
+        "lora_dropout": 0.05,
+        "target_modules": "all-linear",
+        "target_parameters": [
+            "7.mlp.experts.gate_up_proj",
+            "7.mlp.experts.down_proj",
+            "15.mlp.experts.gate_up_proj",
+            "15.mlp.experts.down_proj",
+            "23.mlp.experts.gate_up_proj",
+            "23.mlp.experts.down_proj",
+        ],
+        "bias": "none",
+        "task_type": "CAUSAL_LM",
+    },
+
+    use_quantization=True,
+    quantization_config={
+        "dequantize": True,
+        "load_in_4bit": False,
+        # Optional MXFP4 config is auto-applied by training script if available
+    },
+
+    # ============================================================================
+    # LOGGING & EVAL
+    # ============================================================================
+    eval_strategy="steps",
+    eval_steps=200,
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=500,
+    save_total_limit=3,
+    save_only_model=True,
+    metric_for_best_model="eval_loss",
+    greater_is_better=False,
+    load_best_model_at_end=False,
+    eval_accumulation_steps=2,
+    eval_batch_size=1,
+    eval_ratio=0.01,
+    test_ratio=0.01,
+
+    # ============================================================================
+    # MONITORING & HUB
+    # ============================================================================
+    enable_tracking=True,
+    log_artifacts=False,
+    log_metrics=True,
+    log_config=True,
+    push_to_hub=False,
+    hub_model_id=None,
+    hub_private_repo=False,
+)
+
+# Quick summary for visibility when the config is imported
+print("\n🩺 GPT-OSS Medical o1 SFT Configuration")
+print("=" * 60)
+print(f"📊 Dataset: {config.dataset_name} [{config.dataset_config}] (medical_o1_sft)")
+print(f"📈 Training: {config.num_train_epochs} epoch | batch {config.batch_size} x acc {config.gradient_accumulation_steps}")
+print(f"🧠 LoRA Rank: {config.lora_config['r']}")
+print(f"📏 Sequence Length: {config.max_seq_length}")
+print(f"🎵 Harmony Format: {'Enabled' if config.use_harmony_format else 'Disabled'}")
+print("=" * 60)
+

launch.sh

@@ -267,6 +267,12 @@ show_training_configs() {
         echo "   - Learning Rate: Configurable"
         echo "   - Maximum flexibility with all parameters"
         echo ""
+        echo "8. GPT-OSS Medical o1 SFT (Reasoning)"
+        echo "   - Model: openai/gpt-oss-20b"
+        echo "   - Dataset: FreedomIntelligence/medical-o1-reasoning-SFT"
+        echo "   - Format: Question | Complex_CoT | Response"
+        echo "   - Harmony formatting with optional system/developer messages"
+        echo ""
     fi
 }
 
@@ -376,6 +382,17 @@ get_training_config() {
             MAX_SEQ_LENGTH=1024
             CONFIG_FILE="config/train_gpt_oss_openhermes_fr_memory_optimized.py"
             ;;
+        "GPT-OSS Medical o1 SFT (Reasoning)")
+            MODEL_NAME="openai/gpt-oss-20b"
+            DATASET_NAME="FreedomIntelligence/medical-o1-reasoning-SFT"
+            MAX_EPOCHS=1
+            BATCH_SIZE=2
+            GRADIENT_ACCUMULATION_STEPS=8
+            LEARNING_RATE=2e-4
+            MAX_SEQ_LENGTH=2048
+            CONFIG_FILE="config/train_gpt_oss_medical_o1_sft.py"
+            generate_medical_o1_sft_config
+            ;;
         "GPT-OSS Custom Dataset")
             MODEL_NAME="openai/gpt-oss-20b"
             DATASET_NAME="legmlai/openhermes-fr"  # Will be customizable
@@ -411,10 +428,11 @@ get_custom_dataset_config() {
     echo "1. OpenHermes-FR (prompt + accepted_completion fields)"
     echo "2. Messages format (chat conversations)"
     echo "3. Text format (plain text field)"
-    echo "4. Custom format (specify field names)"
+    echo "4. Medical o1 SFT (Question | Complex_CoT | Response)"
+    echo "5. Custom format (specify field names)"
     echo ""
-    
-    select_option "Select dataset format:" "OpenHermes-FR" "Messages format" "Text format" "Custom format" DATASET_FORMAT
+
+    select_option "Select dataset format:" "OpenHermes-FR" "Messages format" "Text format" "Medical o1 SFT" "Custom format" DATASET_FORMAT
     
     case "$DATASET_FORMAT" in
         "OpenHermes-FR")
@@ -435,6 +453,18 @@ get_custom_dataset_config() {
             DATASET_FORMAT_CODE="text"
             FILTER_BAD_ENTRIES="false"
             ;;
+        "Medical o1 SFT")
+            INPUT_FIELD="Question"
+            TARGET_FIELD="Response"
+            DATASET_FORMAT_CODE="medical_o1_sft"
+            FILTER_BAD_ENTRIES="false"
+            # Field mappings and prefixes
+            get_input "Question field name" "Question" MED_Q_FIELD
+            get_input "Reasoning field name" "Complex_CoT" MED_REASON_FIELD
+            get_input "Response field name" "Response" MED_RESP_FIELD
+            get_input "Reason prefix (before reasoning)" "Reasoning: " MED_REASON_PREFIX
+            get_input "Answer prefix (before final answer)" "Final Answer: " MED_ANSWER_PREFIX
+            ;;
         "Custom format")
             get_input "Input field name" "prompt" INPUT_FIELD
             get_input "Target field name (leave empty if not needed)" "accepted_completion" TARGET_FIELD
@@ -442,6 +472,12 @@ get_custom_dataset_config() {
             get_input "Filter bad entries? (true/false)" "false" FILTER_BAD_ENTRIES
             ;;
     esac
+
+    # Optional Harmony context
+    echo ""
+    print_info "💬 Harmony Context (optional)"
+    get_input "System message" "You are GPT-Tonic, a large language model trained by TonicAI." SYSTEM_MESSAGE
+    get_input "Developer message" "You are an intelligent assistant that can answer customer service queries" DEVELOPER_MESSAGE
     
     # Dataset Filtering Options
     echo ""
@@ -492,6 +528,22 @@ get_custom_dataset_config() {
     update_enhanced_gpt_oss_config
 }
 
+# Function to materialize a default Medical o1 SFT config file
+generate_medical_o1_sft_config() {
+    print_info "Ensuring medical o1 SFT configuration exists..."
+    if [ -f "config/train_gpt_oss_medical_o1_sft.py" ]; then
+        print_status "Medical o1 SFT config already present"
+        return
+    fi
+    cat > config/train_gpt_oss_medical_o1_sft.py << 'EOF'
+"""
+Auto-generated placeholder. A richer version will be imported at runtime.
+"""
+from config.train_gpt_oss_medical_o1_sft import config  # reuse main config
+EOF
+    print_status "Medical o1 SFT config placeholder created"
+}
+
 # Function to get custom configuration
 get_custom_config() {
     print_step "Custom Configuration Setup"
@@ -574,6 +626,18 @@ config = GPTOSSEnhancedCustomConfig(
     min_length=$MIN_LENGTH,
     max_length=$(if [ -n "$MAX_LENGTH" ]; then echo "$MAX_LENGTH"; else echo "None"; fi),
     
+    # Harmony context
+    system_message=$(if [ -n "$SYSTEM_MESSAGE" ]; then printf '%s' "\"$SYSTEM_MESSAGE\""; else echo "None"; fi),
+    developer_message=$(if [ -n "$DEVELOPER_MESSAGE" ]; then printf '%s' "\"$DEVELOPER_MESSAGE\""; else echo "None"; fi),
+    use_harmony_format=True,
+
+    # Medical o1 SFT mapping (ignored unless dataset_format == 'medical_o1_sft')
+    question_field=$(if [ -n "$MED_Q_FIELD" ]; then echo "\"$MED_Q_FIELD\""; else echo "\"Question\""; fi),
+    reasoning_field=$(if [ -n "$MED_REASON_FIELD" ]; then echo "\"$MED_REASON_FIELD\""; else echo "\"Complex_CoT\""; fi),
+    response_field=$(if [ -n "$MED_RESP_FIELD" ]; then echo "\"$MED_RESP_FIELD\""; else echo "\"Response\""; fi),
+    reason_prefix=$(if [ -n "$MED_REASON_PREFIX" ]; then printf '%s' "\"$MED_REASON_PREFIX\""; else echo "\"Reasoning: \""; fi),
+    answer_prefix=$(if [ -n "$MED_ANSWER_PREFIX" ]; then printf '%s' "\"$MED_ANSWER_PREFIX\""; else echo "\"Final Answer: \""; fi),
+    
     # ============================================================================
     # TRAINING HYPERPARAMETERS
     # ============================================================================
@@ -811,6 +875,7 @@ else
         "GPT-OSS OpenHermes-FR (Recommended)" \
         "GPT-OSS OpenHermes-FR Memory Optimized" \
         "GPT-OSS Custom Dataset" \
+        "GPT-OSS Medical o1 SFT (Reasoning)" \
         TRAINING_CONFIG_TYPE
 fi
 

scripts/training/train_gpt_oss.py

@@ -277,31 +277,66 @@ def apply_dataset_filtering(dataset, config):
     
     return dataset
 
-def format_gpt_oss_harmony(prompt, completion, add_eos_token=True):
+def _build_harmony_text(
+    user_content: str,
+    assistant_content: str,
+    add_eos_token: bool = True,
+    system_message: str | None = None,
+    developer_message: str | None = None,
+) -> str:
+    """Compose a Harmony-formatted conversation with optional system/developer messages.
+
+    Structure (training):
+      <|start|>system<|message|>...<|end|> (optional)
+      <|start|>developer<|message|>...<|end|> (optional)
+      <|start|>user<|message|>...<|end|>
+      <|start|>assistant<|channel|>final<|message|>...<|return|>
     """
-    Format data for GPT-OSS Harmony format following the exact template structure.
-    Based on: https://huggingface.co/openai/gpt-oss-20b/raw/main/chat_template.jinja
-    """
-    # GPT-OSS Harmony format structure (exact template compliance)
-    # User message: <|start|>user<|message|>content<|end|>
-    # Assistant message: <|start|>assistant<|channel|>final<|message|>content<|end|> (inference)
-    # Assistant message: <|start|>assistant<|channel|>final<|message|>content<|return|> (training)
-    
-    harmony_text = f"<|start|>user<|message|>{prompt}<|end|><|start|>assistant<|channel|>final<|message|>{completion}"
-    
+    parts: list[str] = []
+    if system_message:
+        parts.append(f"<|start|>system<|message|>{system_message}<|end|>")
+    if developer_message:
+        parts.append(f"<|start|>developer<|message|>{developer_message}<|end|>")
+    parts.append(f"<|start|>user<|message|>{user_content}<|end|>")
+    parts.append(f"<|start|>assistant<|channel|>final<|message|>{assistant_content}")
     if add_eos_token:
-        # Use <|return|> for training as per template specification
-        # This indicates the end of generation in training
-        harmony_text += "<|return|>"
+        parts[-1] += "<|return|>"
     else:
-        # Use <|end|> for inference
-        harmony_text += "<|end|>"
-    
-    return harmony_text
+        parts[-1] += "<|end|>"
+    return "".join(parts)
+
+def format_gpt_oss_harmony(
+    prompt: str,
+    completion: str,
+    add_eos_token: bool = True,
+    system_message: str | None = None,
+    developer_message: str | None = None,
+) -> str:
+    """
+    Format data for GPT-OSS Harmony format following the exact template structure.
+    Spec: `https://huggingface.co/openai/gpt-oss-20b/raw/main/chat_template.jinja`.
+    """
+    return _build_harmony_text(
+        user_content=prompt,
+        assistant_content=completion,
+        add_eos_token=add_eos_token,
+        system_message=system_message,
+        developer_message=developer_message,
+    )
 
-def format_gpt_oss_harmony_prompt(prompt: str) -> str:
-    """Prefix-only Harmony prompt up to assistant content marker for DPO."""
-    return f"<|start|>user<|message|>{prompt}<|end|><|start|>assistant<|channel|>final<|message|>"
+def format_gpt_oss_harmony_prompt(
+    prompt: str,
+    system_message: str | None = None,
+    developer_message: str | None = None,
+) -> str:
+    """Prefix-only Harmony prompt up to assistant content marker for DPO, with optional context."""
+    parts: list[str] = []
+    if system_message:
+        parts.append(f"<|start|>system<|message|>{system_message}<|end|>")
+    if developer_message:
+        parts.append(f"<|start|>developer<|message|>{developer_message}<|end|>")
+    parts.append(f"<|start|>user<|message|>{prompt}<|end|><|start|>assistant<|channel|>final<|message|>")
+    return "".join(parts)
 
 def process_dataset_format(dataset, config):
     """Process dataset based on format configuration with exact GPT-OSS Harmony compliance"""
@@ -321,6 +356,8 @@ def process_dataset_format(dataset, config):
     add_eos_token = getattr(config, 'add_eos_token', True)
     use_harmony_format = getattr(config, 'use_harmony_format', True)
     trainer_type = getattr(config, 'trainer_type', 'sft')
+    system_message = getattr(config, 'system_message', None)
+    developer_message = getattr(config, 'developer_message', None)
     
     print(f"Processing dataset format: {dataset_format}")
     print(f"Input field: {input_field}, Target field: {target_field}")
@@ -338,7 +375,11 @@ def process_dataset_format(dataset, config):
                 chosen_val = example.get('chosen', example.get(chosen_field or 'chosen', ''))
                 rejected_val = example.get('rejected', example.get(rejected_field or 'rejected', ''))
                 if use_harmony_format:
-                    prompt_text = format_gpt_oss_harmony_prompt(prompt_val)
+                    prompt_text = format_gpt_oss_harmony_prompt(
+                        prompt_val,
+                        system_message=system_message,
+                        developer_message=developer_message,
+                    )
                     chosen_text = (chosen_val or '') + ("<|return|>" if add_eos_token else '')
                     rejected_text = (rejected_val or '') + ("<|return|>" if add_eos_token else '')
                     return {"prompt": prompt_text, "chosen": chosen_text, "rejected": rejected_text}
@@ -355,7 +396,11 @@ def process_dataset_format(dataset, config):
                 chosen_val = example.get(chosen_field, '')
                 rejected_val = example.get(rejected_field, '')
                 if use_harmony_format:
-                    prompt_text = format_gpt_oss_harmony_prompt(prompt_val)
+                    prompt_text = format_gpt_oss_harmony_prompt(
+                        prompt_val,
+                        system_message=system_message,
+                        developer_message=developer_message,
+                    )
                     chosen_text = (chosen_val or '') + ("<|return|>" if add_eos_token else '')
                     rejected_text = (rejected_val or '') + ("<|return|>" if add_eos_token else '')
                     return {"prompt": prompt_text, "chosen": chosen_text, "rejected": rejected_text}
@@ -376,7 +421,13 @@ def process_dataset_format(dataset, config):
             if concatenate_fields:
                 if use_harmony_format:
                     # Use exact GPT-OSS Harmony format from template
-                    text = format_gpt_oss_harmony(prompt, completion, add_eos_token)
+                    text = format_gpt_oss_harmony(
+                        prompt,
+                        completion,
+                        add_eos_token,
+                        system_message=system_message,
+                        developer_message=developer_message,
+                    )
                 else:
                     # Fallback to standard format with separator
                     text = prompt + field_separator + completion
@@ -414,7 +465,13 @@ def process_dataset_format(dataset, config):
                 
                 if user_message and assistant_message:
                     # Use GPT-OSS Harmony format
-                    text = format_gpt_oss_harmony(user_message, assistant_message, add_eos_token)
+                    text = format_gpt_oss_harmony(
+                        user_message,
+                        assistant_message,
+                        add_eos_token,
+                        system_message=system_message,
+                        developer_message=developer_message,
+                    )
                 else:
                     # Fallback to simple concatenation
                     text = ""
@@ -438,6 +495,44 @@ def process_dataset_format(dataset, config):
         
         dataset = dataset.map(format_messages, remove_columns=dataset.column_names, num_proc=num_proc)
         
+    elif dataset_format == "medical_o1_sft":
+        # Process Medical-o1 SFT format: Question | Complex_CoT | Response
+        # Defaults align with FreedomIntelligence/medical-o1-reasoning-SFT
+        question_field = getattr(config, 'question_field', input_field or 'Question')
+        reasoning_field = getattr(config, 'reasoning_field', 'Complex_CoT')
+        response_field = getattr(config, 'response_field', target_field or 'Response')
+        reason_prefix = getattr(config, 'reason_prefix', 'Reasoning: ')
+        answer_prefix = getattr(config, 'answer_prefix', 'Final Answer: ')
+
+        def format_medical(example):
+            q = example.get(question_field, '') or ''
+            cot = example.get(reasoning_field, '') or ''
+            ans = example.get(response_field, '') or ''
+
+            # Combine reasoning and final answer in a single assistant turn
+            assistant_text = "\n\n".join(
+                [s for s in [
+                    f"{reason_prefix}{cot}".strip() if cot else '',
+                    f"{answer_prefix}{ans}".strip() if ans else ''
+                ] if s]
+            ) or ans
+
+            if use_harmony_format:
+                text = format_gpt_oss_harmony(
+                    q,
+                    assistant_text,
+                    add_eos_token,
+                    system_message=system_message,
+                    developer_message=developer_message,
+                )
+            else:
+                text = f"Q: {q}\n\n{assistant_text}"
+                if add_eos_token:
+                    text += "</s>"
+            return {"text": text}
+
+        dataset = dataset.map(format_medical, remove_columns=dataset.column_names, num_proc=num_proc)
+
     elif dataset_format == "text":
         # Process plain text format
         text_field = input_field