Pretraining Objective와 Loss 이해

import torch
import torch.nn as nn
from transformers import AutoTokenizer

class SFTCompletionOnlyDataCollator:
    """
    User Prompt 구간을 탐색하여 해당 토큰의 라벨을 -100으로 소거하는 
    상용 레벨의 Completion-only Loss 마스킹 콜레이터 구현체.
    """
    def __init__(self, tokenizer: AutoTokenizer, response_template: str):
        self.tokenizer = tokenizer
        self.response_template = response_template
        # 응답 시작을 알리는 템플릿 문자열을 토큰 ID 패스로 사전에 인코딩
        self.response_token_ids = self.tokenizer.encode(self.response_template, add_special_tokens=False)

    def __call__(self, batch_examples):
        input_ids_list = []
        labels_list = []
        attention_masks_list = []
        
        # 가상의 패딩 차원을 맞추기 위한 최대 길이 계산 (Batch 내 동적 최적화)
        max_len = max(len(ex["input_ids"]) for ex in batch_examples)
        
        for ex in batch_examples:
            w_input_ids = ex["input_ids"][:]
            w_attn_mask = [1] * len(w_input_ids)
            
            # 초기 정답 라벨은 input_ids를 그대로 복사 (Causal LM 학습 기본형)
            w_labels = w_input_ids[:]
            
            # Response Template 토큰 시퀀스가 시작되는 인덱스 검색 (Sublist Matching)
            response_start_idx = -1
            for i in range(len(w_input_ids) - len(self.response_token_ids) + 1):
                if w_input_ids[i : i + len(self.response_token_ids)] == self.response_token_ids:
                    # Template 토큰이 끝난 직후, 즉 실제 '답변 토큰'이 시작되는 위치 지정
                    response_start_idx = i + len(self.response_token_ids)
                    break
            
            if response_start_idx != -1:
                # 0번 인덱스부터 실제 답변 시작 전까지의 모든 프롬프트 토큰 위치를 -100으로 오버라이드
                for k in range(response_start_idx):
                    w_labels[k] = -100
            else:
                # 안전 장치: 템플릿이 검출되지 않은 비정상 데이터는 시퀀스 전체를 마스킹하여 무력화
                w_labels = [-100] * len(w_labels)
                
            # 패딩 처리 (Right Padding 기준)
            pad_len = max_len - len(w_input_ids)
            w_input_ids += [self.tokenizer.pad_token_id] * pad_len
            w_labels += [-100] * pad_len
            w_attn_mask += [0] * pad_len
            
            input_ids_list.append(w_input_ids)
            labels_list.append(w_labels)
            attention_masks_list.append(w_attn_mask)
            
        return {
            "input_ids": torch.tensor(input_ids_list, dtype=torch.long),
            "labels": torch.tensor(labels_list, dtype=torch.long),
            "attention_mask": torch.tensor(attention_masks_list, dtype=torch.long)
        }

# 데이터 파이프라인 목업 검증 스크립트
if __name__ == "__main__":
    from transformers import AutoTokenizer
    # Qwen 토크나이저 바인딩
    tk = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B-Instruct")
    if tk.pad_token is None:
        tk.pad_token = tk.eos_token
        
    # 가상의 대화 원문 구성
    prompt_part = "<|im_start|>user\nPostgreSQL 아키텍처에 대해 설명하라.<|im_end|>\n"
    response_part = "<|im_start|>assistant\nPostgreSQL은 다중 프로세스(Shared Memory) 아키텍처 기반의 RDBMS입니다."
    full_text = prompt_part + response_part
    
    encoded = tk(full_text, add_special_tokens=False)
    sample_batch = [{"input_ids": encoded["input_ids"]}]
    
    # Qwen 대화 템플릿 규격 상, assistant 발화 시작 플래그 문자열을 지표로 설정
    collator = SFTCompletionOnlyDataCollator(tokenizer=tk, response_template="<|im_start|>assistant\n")
    batch = collator(sample_batch)
    
    print("=== Loss Masking 텐서 정밀 प्रो토타입 프로파일링 ===")
    print(f"인코딩된 토큰 시퀀스 전체 길이: {batch['input_ids'].size(1)}")
    
    # 디버깅: 정수 ID가 매핑된 구조 역추적
    for idx, (token_id, label_id) in enumerate(zip(batch["input_ids"][0], batch["labels"][0])):
        decoded_token = tk.decode([token_id]).replace("\n", "\\n")
        print(f"[{idx:02d}] Token: {decoded_token:<25} -> Label Target ID: {label_id.item()}")