I am trying to use a HuggingFace model (Bloom) to generate text and evaluate the token probabilities of all possible completions. I want to generate something alongside OpenAI’s ‘top logprobs’ condition, so we get a dictionary such as: { " No": -1.2539079, " Yes": -0.36010918, " No": -5.115541, " High": -6.792852, " Maybe": -6.834635 } This is the code I have now:
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
torch.set_default_tensor_type(torch.cuda.FloatTensor)
model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", use_cache=True)
tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
input_ids = tokenizer(prompt, return_tensors = "pt").to(0)
outputs = model.generate(**input_ids, max_length = 450, return_dict_in_generate = True, output_scores = True)
print(tokenizer.decode(outputs[0][0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"]))
I am not sure how to approach this – I heard setting “return_dict_in_generate” and “output_scores=True” could do this, but I still do not know how to access those scores.