# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/misc/ml.dataset.lang.ipynb.

# %% auto 0
__all__ = ['alpaca2sharegpt', 'alpaca2openai', 'sharegpt2openai_message', 'openai2sharegpt_message', 'sharegpt2openai',
           'openai2sharegpt', 'partition_input_output_messages']

# %% ../../../nbs/misc/ml.dataset.lang.ipynb 3
def alpaca2sharegpt(example):
    messages = []
    if "instruction" in example:
        messages.append({"from": "system", "value": example["instruction"]})
    messages.append({"from": "human", "value": example["input"]})
    messages.append({"from": "gpt", "value": example["output"]})
    return {"messages": messages}


def alpaca2openai(example):
    messages = []
    if "instruction" in example:
        messages.append({"role": "system", "content": example["instruction"]})
    messages.append({"role": "user", "content": example["input"]})
    messages.append({"role": "assistant", "content": example["output"]})
    return {"messages": messages}


def sharegpt2openai_message(message):
    role = {"system": "system", "human": "user", "gpt": "assistant"}[message["from"]]
    return {"role": role, "content": message["value"]}


def openai2sharegpt_message(message):
    frm = {"system": "system", "user": "human", "assistant": "gpt"}[message["role"]]
    return {"from": frm, "value": message["content"]}


def sharegpt2openai(example):
    return {
        "messages": [sharegpt2openai_message(message) for message in example["messages"]],
    }


def openai2sharegpt(example):
    return {
        "messages": [openai2sharegpt_message(message) for message in example["messages"]],
    }

# %% ../../../nbs/misc/ml.dataset.lang.ipynb 9
def partition_input_output_messages(example: dict) -> dict:
    result = {}
    last_message = example["messages"][-1]
    if last_message['role'] == 'assistant':
        result["input"] = example["messages"][:-1]
        result["output"] = example["messages"][-1:]
    else:
        result["input"] = example["messages"]
    return result
    
