score.py

import torch
from torch import Tensor as _T

from transformers import AutoModelForCausalLM, AutoTokenizer

from torch.nn import functional as F

from typing import List, Optional

from shared_util.generate import full_tokenize


@torch.no_grad()
def score_sequences(
    model: AutoModelForCausalLM,
    input_id_sequences: Optional[_T] = None,    # [batch, sequence length]
    text_sequences: Optional[List[str]] = None,
    tokenizer: Optional[AutoTokenizer] = None,
    device: str = 'cuda'
):
    """
    Given a text sequence or input id sequence, tokenize it and get the sequence of logits for each token
    Because LLMs are next-token predictors, have to offset the logit and token sequences by one
    """
    if input_id_sequences is None:
        model_inputs = full_tokenize(prompts = text_sequences, tokenizer = tokenizer, device = device)
        input_ids = model_inputs['input_ids'].to(device)    # [batch, sequence length]
    else:
        assert text_sequences is None and tokenizer is None
        input_ids = input_id_sequences
    
    # logits[b,l,w] gives logit of input_ids[b,l+1] = w
    # shifted_logits[b,l,w] gives logit of shifted_input_ids[b,l] = w
    outputs = model(input_ids)
    logits = outputs.logits                             # [batch, sequence length, vocab size]
    shifted_logits = logits[:, :-1, :]                  # remove last token prediction (would be a prediction)
    shifted_input_ids = input_ids[:, 1:]                # remove first token (no prediction for it)
    
    # token_log_probs[b,l] = shifted_logits[b,l,shifted_input_ids[b,l]]
    # TODO: could make this more efficient, by not doing the full log_softmax
    log_probs = F.log_softmax(shifted_logits, dim=-1)   # log probabilities [batch, length - 1, vocab]
    token_log_probs = log_probs.gather(2, shifted_input_ids.unsqueeze(-1)).squeeze(-1)  # Actual tokens log probs [batch, length - 1]
    return token_log_probs