(tokenizer, model, inputs: List[str])
| 135 | |
| 136 | |
| 137 | def get_logits(tokenizer, model, inputs: List[str]): |
| 138 | input_ids = tokenizer(inputs, padding=False)["input_ids"] |
| 139 | input_ids = torch.tensor(input_ids, device=model.device) |
| 140 | tokens = {"input_ids": input_ids} |
| 141 | |
| 142 | outputs = model(input_ids)["logits"] |
| 143 | logits = outputs[:, -1, :] |
| 144 | log_probs = torch.nn.functional.softmax(logits, dim=-1) |
| 145 | return log_probs, {"tokens": tokens} |
| 146 | |
| 147 | |
| 148 | @torch.no_grad() |