kf-deberta-gen

from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("solonsophy/kf-deberta-gen")
model = AutoModelForMaskedLM.from_pretrained("solonsophy/kf-deberta-gen")

import torch
import torch.nn.functional as F

def generate_diffusion(model, tokenizer, question, num_steps=15, max_answer_len=80):
    model.eval()
    device = next(model.parameters()).device
    
    MASK_ID = tokenizer.mask_token_id
    CLS_ID = tokenizer.cls_token_id
    SEP_ID = tokenizer.sep_token_id
    
    # 질문 토큰화
    q_tokens = tokenizer.encode(question, add_special_tokens=False)[:100]
    
    # 초기: [CLS] Q [SEP] [MASK]*N
    input_ids = [CLS_ID] + q_tokens + [SEP_ID] + [MASK_ID] * max_answer_len
    input_ids = torch.tensor([input_ids[:256]], device=device)
    answer_start = len(q_tokens) + 2
    
    # Iterative denoising
    for step in range(num_steps):
        with torch.no_grad():
            logits = model(input_ids).logits
        
        mask_pos = (input_ids[0, answer_start:] == MASK_ID).nonzero().squeeze(-1) + answer_start
        if len(mask_pos) == 0:
            break
        
        # Confidence 기반 unmask
        mask_logits = logits[0, mask_pos] / 0.8  # temperature
        probs = F.softmax(mask_logits, dim=-1)
        tokens = torch.multinomial(probs, 1).squeeze(-1)
        conf = probs.gather(1, tokens.unsqueeze(-1)).squeeze(-1)
        
        k = max(1, len(mask_pos) // (num_steps - step))
        top_idx = conf.topk(k).indices
        input_ids[0, mask_pos[top_idx]] = tokens[top_idx]
    
    # 결과 추출
    answer = input_ids[0, answer_start:]
    answer = answer[(answer != MASK_ID) & (answer != tokenizer.pad_token_id)]
    return tokenizer.decode(answer, skip_special_tokens=True)

# 사용 예시
answer = generate_diffusion(model, tokenizer, "인공지능이란 무엇인가요?")
print(answer)