convert_for_training.py

"""
Convert dataset.jsonl to various fine-tuning formats

Supported formats:
1. Instruction-following (for instruction-tuned models)
2. Conversational (for chat models)
3. Completion (for base models)
"""
import json
from pathlib import Path
from typing import List, Dict, Any


def format_structured_response(response: Dict[str, Any]) -> str:
    """Format the structured response into readable text"""
    parts = []
    
    if response.get("facts"):
        parts.append("FACTS:")
        for fact in response["facts"]:
            parts.append(f"  • {fact}")
    
    if response.get("uncertainties"):
        parts.append("\nUNCERTAINTIES:")
        for unc in response["uncertainties"]:
            parts.append(f"  • {unc}")
    
    if response.get("analysis"):
        parts.append("\nANALYSIS:")
        for analysis in response["analysis"]:
            parts.append(f"  • {analysis}")
    
    if response.get("guidance"):
        parts.append("\nGUIDANCE:")
        for guidance in response["guidance"]:
            parts.append(f"  • {guidance}")
    
    if "confidence" in response:
        parts.append(f"\nConfidence: {response['confidence']:.2f}")
    
    return "\n".join(parts)


def convert_to_instruction_format(sample: Dict[str, Any], role: str = "civilian") -> Dict[str, str]:
    """Convert to instruction-following format (instruction + response)"""
    scenario = sample["scenario"]
    category = sample["category"]
    response = sample["responses"][role]
    
    instruction = f"""You are a crisis response expert. Analyze the following crisis scenario and provide a structured response.

Category: {category}

Scenario:
{scenario}

Provide your analysis with:
- Key facts you can observe
- Uncertainties that need clarification
- Your analysis of the situation
- Actionable guidance

Role: {role.title()}"""
    
    output = format_structured_response(response)
    
    return {
        "instruction": instruction,
        "output": output,
        "category": category,
        "role": role
    }


def convert_to_conversational_format(sample: Dict[str, Any], role: str = "civilian") -> Dict[str, Any]:
    """Convert to conversational format (messages with roles)"""
    scenario = sample["scenario"]
    category = sample["category"]
    response = sample["responses"][role]
    
    messages = [
        {
            "role": "system",
            "content": f"You are a crisis response expert specializing in {category} scenarios. Provide structured, actionable guidance based on crisis situations."
        },
        {
            "role": "user",
            "content": f"Category: {category}\n\nScenario:\n{scenario}\n\nAnalyze this crisis situation and provide guidance. Role: {role.title()}"
        },
        {
            "role": "assistant",
            "content": format_structured_response(response)
        }
    ]
    
    return {
        "messages": messages,
        "category": category,
        "role": role
    }


def convert_to_completion_format(sample: Dict[str, Any], role: str = "civilian") -> Dict[str, str]:
    """Convert to completion format (prompt + completion)"""
    scenario = sample["scenario"]
    category = sample["category"]
    response = sample["responses"][role]
    
    prompt = f"""Category: {category}

Scenario:
{scenario}

As a {role}, analyze this crisis and provide structured guidance:

"""
    
    completion = format_structured_response(response)
    
    return {
        "prompt": prompt,
        "completion": completion,
        "category": category,
        "role": role
    }


def convert_dataset(
    input_path: Path,
    output_path: Path,
    format_type: str = "instruction",
    role: str = "civilian",
    include_both_roles: bool = False
):
    """
    Convert dataset to training format
    
    Args:
        input_path: Path to dataset.jsonl
        output_path: Path to output file
        format_type: "instruction", "conversational", or "completion"
        role: "civilian" or "first responder"
        include_both_roles: If True, create separate examples for each role
    """
    samples = []
    
    # Load samples
    with open(input_path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                samples.append(json.loads(line))
    
    converted = []
    
    for sample in samples:
        roles_to_process = ["civilian", "first responder"] if include_both_roles else [role]
        
        for r in roles_to_process:
            if format_type == "instruction":
                converted.append(convert_to_instruction_format(sample, r))
            elif format_type == "conversational":
                converted.append(convert_to_conversational_format(sample, r))
            elif format_type == "completion":
                converted.append(convert_to_completion_format(sample, r))
            else:
                raise ValueError(f"Unknown format: {format_type}")
    
    # Write output
    with open(output_path, "w", encoding="utf-8") as f:
        for item in converted:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
    
    print(f"Converted {len(samples)} samples to {len(converted)} training examples")
    print(f"Format: {format_type}")
    print(f"Output: {output_path}")


if __name__ == "__main__":
    import typer
    
    app = typer.Typer(help="Convert dataset for fine-tuning")
    
    @app.command()
    def convert(
        input_file: str = typer.Option("data/dataset.jsonl", help="Input dataset file"),
        output_file: str = typer.Option("data/training_data.jsonl", help="Output training file"),
        format_type: str = typer.Option("instruction", help="Format: instruction, conversational, or completion"),
        role: str = typer.Option("civilian", help="Role: civilian or first_responder"),
        both_roles: bool = typer.Option(False, help="Include both roles as separate examples")
    ):
        """Convert dataset to fine-tuning format"""
        convert_dataset(
            Path(input_file),
            Path(output_file),
            format_type=format_type,
            role=role.replace("_", " "),
            include_both_roles=both_roles
        )
    
    if __name__ == "__main__":
        app()