잡소리:
아
추워

개념적인 부분은 이전 글을 참조하자

Encoder 코드 살펴보기

에서 우리가 볼 것은 dpr/models/biencoder.py의 Biencoder 클래스와 dpr/models/hf_models.py이다. 코드를 다 적기엔 길기에 옆에 띄워두고 같이 보도록 하자

Biencoder.py

init

모델이 question(Question encoder)과 ctx(passage encoder) 두개가 있는 것을 확인할 수 있다.
fix_...는 해당 모델을 고정할지 여부이다.

get_representation

Encoder 1개를 호출하고, 해당 encoder에 대한 출력값 3개를 반환한다.
이때 만약 fix된 모델이라면 nograd로 가중치를 업데이트하지 않는다. 다만 encoder이후의 projection layer, loss는 requires_grad를 통해서 학습되도록 한다.
return

• sequence_output(토큰별 hidden state. [batch, seq_len, hidden]
• pooled_output(문장 요약 벡터([CLS] 토큰의 hidden state), [batch_size, hidden_dim]
• hidden_states(레이어별 히든 스택)

forward

Question vector와 Passage vector를 각각 뽑아내서 반환
단, 이때 q_pooled_out과 ctx_pooled_out만을 사용하면서 문장 전체에 대한 대표 벡터만 리턴

create_biencoder_input

해당 함수가 핵심
Objective: 학습용 batch 만들기. Question tensor은 [B, Lq], Passage tensor은 [Nctx, Lc], 각 질문의 positive 문서가 ctx_tensor의 어디 위치인지도 기록하는것이 목표

1) sample마다 positivie 1개 + negative 뽑기

positive_ctx = sample.positivie_passages[0] # shuffle_positives면 랜덤으로 1개 추출
neg_ctxs = sample_negative
hard_neg_ctxs = sample.hard_negative_passages

이때 만약 hard_neg_fallback = True이고 hard negative가 부족하다면 일반 negative에서 hard_negative 개수만큼 대체함

if hard_neg_fallback and len(hard_neg_ctxs) == 0
	hard_neg_ctxs = neg_ctxs[0:num_hard_negatives]

이를 원하는 개수(num_other_negatives or num_neg_ctxs)만큼 잘라서 최종 문서에선 all_ctxs = [positivie_ctx] + neg_ctxs + hard_neg_ctxs

2) ctx를 텐서로 바꾸고 정답 인덱스 기록

sample_ctxs_tensors = [
	tensorizer.text_to_tensor(ctx.text, title=ctx.title if (insert_title and ctx.title) else None)
	for ctx in all_ctxs
]

ctx_tensors.extend(sample_ctxs_tensors)

tokenization + padding으로 fixed-size tensor을 만들어줌

 	positive_ctx_indices.append(current_ctxs_len)
	hard_neg_ctx_indices.append(...)

positive_ctx가 전체 ctx_tensor에서 시작한 위치가 positive index. 추가로 hard negative의 index도 기록

3) Question(Query) tensor화
query_token 옵션이 있다면 질문 앞에 특정 token을 붙임
특수 토큰([START_ENT])인 경우 span을 골라 넣음(특수 목적용 -> if query_token == "[START_END]"부분은 무시해도 된다.)

if query_token:
	# TODO: tmp workaround for EL, remove or revise
	if query_token == "[START_ENT]":
		query_span = _select_span_with_token(question, tensorizer, token_str=query_token)
		question_tensors.append(query_span)
	else:
		question_tensors.append(tensorizer.text_to_tensor(" ".join([query_token, question])))
else:
	question_tensors.append(tensorizer.text_to_tensor(question))

4) 최종 텐서 구성
passage tensor

ctxs_tensor = torch.cat([ctx.view(1, -1) for ctx in ctx_tensors], dim=0)

Question tensor

questions_tensor = torch.cat([q.view(1, -1) for q in question_tensors], dim=0)

class BiEncoderNllLoss

얘는 위의 함수들처럼 Biencoder class가 아닌 다른클래스인데, 간단하게 question 벡터와 ctx벡터에서 score를 계산한다. 이후 해당 score를 softmax하고, 이런 softmax_score를 바탕으로 loss를 계산한다는 softmax cross entrophy의 교과서적인 방법이다.
약간의 차이점으로는 class가 일반적인 라벨이 아닌 배치 내의 passage라는 것이다. 이 부분이 논문에서 말한 in-batch negatives이다.

hf_models.py

get_bert_biencoder_components

DPR의 Dual Encoder를 조립하는 부분이다.

question_encoder = HFBertEncoder.init_encoder(
	cfg.encoder.pretrained_model_cfg,
	projection_dim=cfg.encoder.projection_dim,
	dropout=dropout,
	pretrained=cfg.encoder.pretrained,
	**kwargs
    )
ctx_encoder = HFBertEncoder.init_encoder(
	cfg.encoder.pretrained_model_cfg,
	projection_dim=cfg.encoder.projection_dim,
	dropout=dropout,
	pretrained=cfg.encoder.pretrained,
	**kwargs
)

크게 특별한 부분은 없지만, question과 ctx가 모두 BERT 기반이지만 가중치가 별도로 구성되어있다. 그리고 이런 두 BERT를

biencoder = BiEncoder(question_encoder, ctx_encoder, fix_ctx_encoder=fix_ctx_encoder)

합친다. 이 부분이 DPR에서 말한 아이디어처럼 $E_q$와 $E_p$를 묶은 것이다.
그 밑에는 평범하게 AdamW로 optimizer하고, tokenziar하고 한다.

get_bert_reader_components

DPR의 두 단계에서 말한 부분 중 reader에 속하는 부분으로, 위에서 뽑은 top-k를 바탕으로 가져온 passage 중 정답 span을 뽑는 부분이다.

encoder = HFBertEncoder.init_encoder(
	cfg.encoder.pretrained_model_cfg,
	projection_dim=cfg.encoder.projection_dim,
	dropout=dropout,
	pretrained=cfg.encoder.pretrained,
	**kwargs
)

BERT encoder를 만들고

reader = Reader(encoder, hidden_size)

Span 추출용 헤드를 얹는다.
이후 위와 동일하게 optimizer하고 tensorizer한다.

Class HFBertEncoder

Huggingface에 있는 BertModel을 inherited(상속)하고 DPR에 필요한 형태로 적절하게 변형한다.(wrap한다고 한다.)

__init__

BERT model 초기화, projection차원 설정, 가중치 초기화

init_encoder

cfg = BertConfig.from_pretrained(cfg_name if cfg_name else "bert-base-uncased")

pre-training된 BERT를 그대로 가져오고

if dropout != 0:
	cfg.attention_probs_dropout_prob = dropout
	cfg.hidden_dropout_prob = dropout

그 위에 원하는 dropout을 덮어쓴다.

if pretrained:
	return cls.from_pretrained(cfg_name, config=cfg, project_dim=projection_dim, **kwargs)
else:
	return HFBertEncoder(cfg, project_dim=projection_dim)

이부분은 pre-training BERT와 scratch BERT 둘 다 실험하기 위한 부분이다.

forward

out = super().forward(
	input_ids=input_ids,
	token_type_ids=token_type_ids,
	attention_mask=attention_mask,
)

순수 BERT를 그대로 실행하고

sequence_output = out.last_hidden_state
pooled_output = None
hidden_states = out.hidden_states

if문은 Hugging face의 버전에 따라서 맞춰둔 설정이라 넘어가면, 해당 부분에서 sequence_output을 얻고

if isinstance(representation_token_pos, int):
	pooled_output = sequence_output[:, representation_token_pos, :]

위의 sequence_output을 바탕으로 DPR식으로 pooled representation을 생성한다.
해당 pooled_output을 바탕으로 유사도 계산을 진행한다.

결론

DRP의 코드를 아주 간단하게만 살펴보았다. 아마 아예 배경지식이 없는 경우(BERT도 모른다거나) 이해가 안될텐데, 그건 다른 글들을 찾아보길 바란다. 결론적으로 설명된 것 처럼 batch를 만들어내고, $E_q$와 $E_p$ 두 BERT를 묶어서 둘 다 학습시킨다는 점이 중요하다.