실습 2 — 대화 형식으로 바꾸고 나누기

# 첫 행을 Prompt-Completion 대화 형식으로 변환
# loc[행, 열]: 0번 행의 prompt 칸 값
first_prompt = persona_df.loc[0, "prompt"]  # 첫 번째 질문 문자열입니다.
first_response = persona_df.loc[0, "response"]  # 첫 번째 답변 문자열입니다.
first_record = {  # TRL이 읽을 대화 예시 하나를 만듭니다.
    # prompt: role 이 user 인 메시지 하나를 담은 리스트
    "prompt": [{"role": "user", "content": first_prompt}],  # 사용자 메시지입니다.
    # completion: role 이 assistant 인 메시지 하나를 담은 리스트(모델이 배울 답)
    "completion": [{"role": "assistant", "content": first_response}],  # 정답 assistant 메시지입니다.
}
# json.dumps(사전, ensure_ascii=False, indent=2): 한글을 그대로 두고 두 칸 들여쓰기로 보기 좋게 찍는다
# 결과: "prompt" 아래에 user 메시지, "completion" 아래에 assistant 메시지가 든 JSON
print(json.dumps(first_record, ensure_ascii=False, indent=2))

records = [  # 모든 행을 같은 대화 구조로 바꿉니다.
    {"prompt": [{"role": "user", "content": prompt}], "completion": [{"role": "assistant", "content": response}]}  # 질문과 답변을 역할별로 담습니다.
    # zip(질문 열, 답변 열): 같은 행의 질문과 답을 짝지어 하나씩 꺼낸다
    for prompt, response in zip(persona_df["prompt"], persona_df["response"])  # 두 열을 같은 순서로 순회합니다.
]
# 결과: 변환한 예시 수: 327
print("변환한 예시 수:", len(records))

# 학습 80%, 검증 10%, 평가 10%로 분할
# sample(frac=1): 전체(100%)를 무작위 순서로 뽑기 = 섞기
# random_state=42: 섞는 순서를 고정하는 씨앗. reset_index(drop=True): 번호를 0부터 다시 매기고 옛 번호는 버린다
shuffled = persona_df.sample(frac=1, random_state=42).reset_index(drop=True)  # 재현 가능한 순서로 모든 행을 섞습니다.
# int(327 * 0.1) = int(32.7) = 32. 두 변수에 같은 값 32 를 넣는다
n_test = n_valid = int(len(shuffled) * 0.1)  # 검증·평가 데이터 수를 각각 10%로 정합니다.
splits = {  # 세 분할의 행 범위를 딕셔너리로 정리합니다.
    # iloc[:32]: 0~31번 행
    "test": shuffled.iloc[:n_test],  # 앞 10%는 최종 평가에 남깁니다.
    # iloc[32:64]: 32~63번 행
    "valid": shuffled.iloc[n_test:n_test + n_valid],  # 다음 10%는 학습 중 검증에 씁니다.
    # iloc[64:]: 64번부터 끝까지 263행
    "train": shuffled.iloc[n_test + n_valid:],  # 나머지 80%는 학습에 씁니다.
}

# 분할별 표를 Hugging Face Dataset으로 변환
# Dataset = 표 하나, DatasetDict = 이름 붙인 Dataset 여러 개를 담은 사전
from datasets import Dataset, DatasetDict  # 표를 학습에 쓰는 Dataset 객체로 바꿉니다.

def to_dataset(df):  # 표 한 조각을 Prompt-Completion Dataset으로 바꿉니다.
    records = [  # 질문과 답변을 역할별로 담습니다.
        {"prompt": [{"role": "user", "content": prompt}], "completion": [{"role": "assistant", "content": response}]}
        for prompt, response in zip(df["prompt"], df["response"])  # 두 열을 같은 순서로 순회합니다.
    ]
    # from_list = 사전의 리스트에서 만들기. 사전의 키(prompt, completion)가 열 이름이 된다
    return Dataset.from_list(records)  # 리스트를 Dataset으로 바꿉니다.


# 사전 내포: splits 의 test, valid, train 마다 to_dataset 을 부른다
persona_ds = DatasetDict({name: to_dataset(part) for name, part in splits.items()})  # 분할마다 Dataset을 만듭니다.
# 결과: test 32행, valid 32행, train 263행. 열(features)은 ['prompt', 'completion']
print(persona_ds)
# 결과: train 의 첫 예시. "주말에 영화나 보러 갈까 하는데 어때?" 와 냉소적인 답
print(persona_ds["train"][0])