(tokenizer, model, inputs: List[str])
| 114 | |
| 115 | |
| 116 | def get_logits(tokenizer, model, inputs: List[str]): |
| 117 | input_ids = tokenizer(inputs, padding=False)["input_ids"] |
| 118 | input_ids = torch.tensor(input_ids, device=model.device) |
| 119 | tokens = {"input_ids": input_ids} |
| 120 | |
| 121 | outputs = model(input_ids)["logits"] |
| 122 | logits = outputs[:, -1, :] |
| 123 | log_probs = torch.nn.functional.softmax(logits, dim=-1) |
| 124 | return log_probs, {"tokens": tokens} |
| 125 | |
| 126 | |
| 127 | @torch.no_grad() |