ReContext: Recursive Evidence Replay as LLM Harness for Long-Context Reasoning

# ReContext 핵심 로직 의사코드 (Hugging Face 기반)
import torch

def get_relevance_scores(model, input_ids, question_token_ids, layer_heads=None):
    """
    모델 내부 attention으로 context 토큰별 relevance 점수 계산
    """
    with torch.no_grad():
        outputs = model(input_ids, output_attentions=True)
    
    # 질문 suffix 토큰(마지막 w=8개)의 attention을 cue로 사용
    cue_positions = list(range(input_ids.shape[1] - 8, input_ids.shape[1]))
    
    relevance_scores = torch.zeros(input_ids.shape[1])
    decay = 0.75
    running_score = None
    
    for layer_idx, attn in enumerate(outputs.attentions):
        # attn shape: (batch, heads, seq, seq)
        cue_attn = attn[0, :, cue_positions, :].mean(dim=(0, 1))  # (seq,)
        if running_score is None:
            running_score = cue_attn
        else:
            running_score = cue_attn + decay * running_score
            running_score = running_score / running_score.sum()
    
    return running_score

def recontext_inference(model, tokenizer, context, question, R=2, K=16):
    """
    ReContext: 재귀적 증거 replay 추론
    """
    # 문장 단위 분리
    sentences = context.split('. ')
    sentence_ranges = []  # 각 문장의 토큰 범위 추적
    
    evidence_pool = []
    current_prompt = context + '\n\n' + question
    
    for round_idx in range(R):
        # 현재 프롬프트 토크나이징
        input_ids = tokenizer(current_prompt, return_tensors='pt').input_ids
        
        # Relevance 점수 계산 (attention 기반)
        scores = get_relevance_scores(model, input_ids, question)
        
        # 원본 컨텍스트 토큰 범위에서만 Top-K 선택
        context_len = len(tokenizer(context).input_ids)
        context_scores = scores[:context_len]
        top_k_indices = context_scores.topk(K).indices
        
        # 선택된 토큰 → 문장으로 매핑
        selected_tokens = tokenizer.convert_ids_to_tokens(
            input_ids[0, top_k_indices]
        )
        new_evidence = extract_sentences_by_tokens(sentences, top_k_indices)
        
        # 중복 제거 후 evidence pool에 추가
        for sent in new_evidence:
            if sent not in evidence_pool:
                evidence_pool.append(sent)
        
        # Replay 프롬프트 재구성: [원본 컨텍스트; 증거; 질문]
        evidence_text = '\n'.join(evidence_pool)
        current_prompt = f"{context}\n\n[Evidence]\n{evidence_text}\n\n{question}"
    
    # 최종 답변 생성
    input_ids = tokenizer(current_prompt, return_tensors='pt').input_ids
    output = model.generate(input_ids, max_new_tokens=100)
    return tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)

# 사용 예시
# answer = recontext_inference(model, tokenizer, long_context, question, R=2, K=16)