Spaces:

wjsdbqls
/

capstone-ai

Sleeping

App Files Files Community

capstone-ai / models /classifier.py

wjsdbqls

Claude Opus 5.5

refactor: KoBERT 이름을 정리하고 미사용 리랭커 경로 제거

1f204c9 6 days ago

Raw History Blame Contribute Delete

4.3 kB

"""
KLUE BERT 기반 문의 카테고리 분류 모델.
구조:
  - klue/bert-base → [CLS] 토큰 → Dropout → Linear → 카테고리 확률
  - 카테고리: 수강신청(0), 성적(1), 졸업(2), 장학금(3), 휴복학(4),
             등록금(5), 기숙사(6), 공결_출석(7), 증명서(8), 기타(9)
attention 구현:
  최신 transformers의 기본 SDPA 구현은 attention 가중치를 반환하지 않으므로
  하이라이팅을 위해 eager 구현을 명시한다.
"""

import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer, AutoTokenizer

# 카테고리 매핑 (generate_qa.py의 CATEGORIES와 동일하게 유지)
CATEGORY_LABELS = [
    "수강신청", "성적", "졸업", "장학금", "휴복학",
    "등록금", "기숙사", "공결_출석", "증명서", "기타"
]
NUM_CLASSES = len(CATEGORY_LABELS)
MODEL_NAME = "klue/bert-base"


class CategoryClassifier(nn.Module):
    """KLUE BERT 파인튜닝 분류 모델."""

    def __init__(self, num_classes: int = NUM_CLASSES, dropout_prob: float = 0.3):
        super().__init__()
        # eager attention: output_attentions 지원을 위해 명시
        self.bert = BertModel.from_pretrained(
            MODEL_NAME,
            attn_implementation="eager",
        )
        hidden_size = self.bert.config.hidden_size
        self.dropout = nn.Dropout(dropout_prob)
        self.classifier = nn.Linear(hidden_size, num_classes)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        token_type_ids: torch.Tensor | None = None,
    ) -> torch.Tensor:
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            output_attentions=True,  # 하이라이팅을 위해 attention 반환
        )
        # [CLS] 토큰의 pooled output 사용
        pooled = outputs.pooler_output          # (batch, hidden)
        pooled = self.dropout(pooled)
        logits = self.classifier(pooled)        # (batch, num_classes)

        # attentions: tuple of (batch, heads, seq, seq) — 마지막 레이어만 반환
        self._last_attentions = outputs.attentions  # inference 시 하이라이팅에 사용
        return logits

    def get_attention_weights(self) -> torch.Tensor:
        """마지막 forward에서의 마지막 레이어 attention 반환. highlight.py에서 사용."""
        if not hasattr(self, "_last_attentions"):
            raise RuntimeError("forward()를 먼저 호출하세요.")
        # 마지막 레이어, 모든 헤드 평균 → (batch, seq, seq)
        last_layer_attn = self._last_attentions[-1]   # (batch, heads, seq, seq)
        return last_layer_attn.mean(dim=1)            # (batch, seq, seq)


def load_tokenizer() -> AutoTokenizer:
    return AutoTokenizer.from_pretrained(MODEL_NAME)


def load_model(checkpoint_path: str | None = None, device: str = "cpu") -> CategoryClassifier:
    """
    모델 로드. checkpoint_path가 있으면 파인튜닝된 가중치 로드.
    """
    model = CategoryClassifier()
    if checkpoint_path:
        state = torch.load(checkpoint_path, map_location=device)
        model.load_state_dict(state)
    model.to(device)
    return model


class InquiryDataset(torch.utils.data.Dataset):
    """학습/검증 데이터셋."""

    def __init__(self, data: list[dict], tokenizer: BertTokenizer, max_length: int = 128):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.texts = [item["question"] for item in data]
        self.labels = [item["label"] for item in data]

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        return {
            "input_ids": encoding["input_ids"].squeeze(0),
            "attention_mask": encoding["attention_mask"].squeeze(0),
            "token_type_ids": encoding.get("token_type_ids", torch.zeros(self.max_length, dtype=torch.long)).squeeze(0),
            "label": torch.tensor(self.labels[idx], dtype=torch.long),
        }

	"""
	KLUE BERT 기반 문의 카테고리 분류 모델.

	구조:
	- klue/bert-base → [CLS] 토큰 → Dropout → Linear → 카테고리 확률
	- 카테고리: 수강신청(0), 성적(1), 졸업(2), 장학금(3), 휴복학(4),
	등록금(5), 기숙사(6), 공결_출석(7), 증명서(8), 기타(9)

	attention 구현:
	최신 transformers의 기본 SDPA 구현은 attention 가중치를 반환하지 않으므로
	하이라이팅을 위해 eager 구현을 명시한다.
	"""

	import torch
	import torch.nn as nn
	from transformers import BertModel, BertTokenizer, AutoTokenizer

	# 카테고리 매핑 (generate_qa.py의 CATEGORIES와 동일하게 유지)
	CATEGORY_LABELS = [
	"수강신청", "성적", "졸업", "장학금", "휴복학",
	"등록금", "기숙사", "공결_출석", "증명서", "기타"
	]
	NUM_CLASSES = len(CATEGORY_LABELS)
	MODEL_NAME = "klue/bert-base"


	class CategoryClassifier(nn.Module):
	"""KLUE BERT 파인튜닝 분류 모델."""

	def __init__(self, num_classes: int = NUM_CLASSES, dropout_prob: float = 0.3):
	super().__init__()
	# eager attention: output_attentions 지원을 위해 명시
	self.bert = BertModel.from_pretrained(
	MODEL_NAME,
	attn_implementation="eager",
	)
	hidden_size = self.bert.config.hidden_size
	self.dropout = nn.Dropout(dropout_prob)
	self.classifier = nn.Linear(hidden_size, num_classes)

	def forward(
	self,
	input_ids: torch.Tensor,
	attention_mask: torch.Tensor,
	token_type_ids: torch.Tensor \| None = None,
	) -> torch.Tensor:
	outputs = self.bert(
	input_ids=input_ids,
	attention_mask=attention_mask,
	token_type_ids=token_type_ids,
	output_attentions=True, # 하이라이팅을 위해 attention 반환
	)
	# [CLS] 토큰의 pooled output 사용
	pooled = outputs.pooler_output # (batch, hidden)
	pooled = self.dropout(pooled)
	logits = self.classifier(pooled) # (batch, num_classes)

	# attentions: tuple of (batch, heads, seq, seq) — 마지막 레이어만 반환
	self._last_attentions = outputs.attentions # inference 시 하이라이팅에 사용
	return logits

	def get_attention_weights(self) -> torch.Tensor:
	"""마지막 forward에서의 마지막 레이어 attention 반환. highlight.py에서 사용."""
	if not hasattr(self, "_last_attentions"):
	raise RuntimeError("forward()를 먼저 호출하세요.")
	# 마지막 레이어, 모든 헤드 평균 → (batch, seq, seq)
	last_layer_attn = self._last_attentions[-1] # (batch, heads, seq, seq)
	return last_layer_attn.mean(dim=1) # (batch, seq, seq)


	def load_tokenizer() -> AutoTokenizer:
	return AutoTokenizer.from_pretrained(MODEL_NAME)


	def load_model(checkpoint_path: str \| None = None, device: str = "cpu") -> CategoryClassifier:
	"""
	모델 로드. checkpoint_path가 있으면 파인튜닝된 가중치 로드.
	"""
	model = CategoryClassifier()
	if checkpoint_path:
	state = torch.load(checkpoint_path, map_location=device)
	model.load_state_dict(state)
	model.to(device)
	return model


	class InquiryDataset(torch.utils.data.Dataset):
	"""학습/검증 데이터셋."""

	def __init__(self, data: list[dict], tokenizer: BertTokenizer, max_length: int = 128):
	self.tokenizer = tokenizer
	self.max_length = max_length
	self.texts = [item["question"] for item in data]
	self.labels = [item["label"] for item in data]

	def __len__(self):
	return len(self.texts)

	def __getitem__(self, idx):
	encoding = self.tokenizer(
	self.texts[idx],
	max_length=self.max_length,
	padding="max_length",
	truncation=True,
	return_tensors="pt",
	)
	return {
	"input_ids": encoding["input_ids"].squeeze(0),
	"attention_mask": encoding["attention_mask"].squeeze(0),
	"token_type_ids": encoding.get("token_type_ids", torch.zeros(self.max_length, dtype=torch.long)).squeeze(0),
	"label": torch.tensor(self.labels[idx], dtype=torch.long),
	}