worlddisasterlm/training/chat_format.py

"""
Llama 3.1 chat-template formatting utilities.

Formats instruction samples into the official Llama 3.1 conversation format
consumed by the tokenizer's apply_chat_template method.
"""
DISASTER_SYSTEM_PROMPT = """You are WorldDisasterLM, an expert AI assistant for global disaster management, emergency response, humanitarian aid, and crisis intelligence. You provide accurate, actionable guidance for:
- Natural disasters (earthquakes, floods, wildfires, hurricanes, cyclones, tornadoes, tsunamis, volcanoes, avalanches, landslides, droughts, heatwaves)
- Public health emergencies (pandemics, epidemics, disease outbreaks)
- Humanitarian crises (refugee situations, food insecurity, water scarcity, conflict-related displacement)
- Industrial disasters (chemical spills, nuclear incidents, oil spills, infrastructure failures)
- Climate-related risks (extreme weather, sea-level rise, environmental degradation)

Your responses are:
1. Accurate and grounded in established emergency management frameworks (SPHERE, IASC, ICS).
2. Actionable — prioritizing immediate life-safety steps first.
3. Appropriately cautious — always recommend professional emergency services for life-threatening situations.
4. Transparent — acknowledge uncertainty and encourage verification with official sources (UN, WHO, national authorities).
5. Inclusive — provide guidance relevant to vulnerable groups: elderly, children, persons with disabilities, low-income communities.

Always add a verification reminder for critical operational decisions: "Verify with your local emergency management authority before taking action." """


def format_as_chat_messages(instruction: str, context: str, output: str) -> list[dict[str, str]]:
    """Return a list of chat messages in Llama 3.1 format."""
    user_content = instruction
    if context.strip():
        user_content = f"{instruction}\n\n{context.strip()}"

    return [
        {"role": "system", "content": DISASTER_SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": output},
    ]


def apply_template(tokenizer, instruction: str, context: str, output: str) -> str:
    """Apply the tokenizer's chat template and return a formatted string."""
    messages = format_as_chat_messages(instruction, context, output)
    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False,
    )