# SillokBert-Scratch 프로젝트 2단계: 모델 아키텍처 정의
# -----------------------------------------------------------------
# 1단계에서 생성된 토크나이저를 기반으로 사전 학습되지 않은(from scratch)
# BERT 모델의 뼈대를 정의하고 정보를 확인합니다.
# -----------------------------------------------------------------
from pathlib import Path
from transformers import BertConfig, BertForMaskedLM, PreTrainedTokenizerFast

def define_sillok_bert_architecture():
    """
    사전 학습된 가중치 없이, BERT-base 구조와 커스텀 BPE 토크나이저 어휘집에
    맞춰 초기화된 '빈' BERT 모델을 정의하고 정보를 출력합니다.
    """
    # --- 경로 설정 ---
    project_dir = Path("/home/work/baro/sillok/sillok_scratch_20250626")
    tokenizer_dir = project_dir / "sillok_tokenizer_bpe_preprocessed"
    tokenizer_file = tokenizer_dir / "tokenizer.json"

    print("--- 2. SillokBert Model Architecture Definition ---")

    # --- 토크나이저 로드 및 PAD 토큰 설정 ---
    tokenizer = PreTrainedTokenizerFast(tokenizer_file=str(tokenizer_file))
    if tokenizer.pad_token is None:
        tokenizer.pad_token = '[PAD]'
    vocab_size = tokenizer.vocab_size

    print(f"로드된 어휘집 크기 (vocab_size): {vocab_size}")
    print(f"PAD 토큰 ID: {tokenizer.pad_token_id}")

    # --- 모델 구성(Configuration) 정의 ---
    config = BertConfig(
        vocab_size=vocab_size,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        max_position_embeddings=512,
        pad_token_id=tokenizer.pad_token_id,
    )
    print("\n생성된 BERT 모델 설정 (Configuration):")
    print(config)

    # --- 모델 생성 및 정보 확인 ---
    model = BertForMaskedLM(config=config)
    print("\n✅ 'From Scratch' BERT 모델이 성공적으로 생성되었습니다.")
    print(f"모델의 총 파라미터 수: {model.num_parameters():,}")

if __name__ == "__main__":
    define_sillok_bert_architecture()