Spaces:

Ilaa-1505
/

llm-visualizer

Running

App Files Files Community

llm-visualizer / modules /generation.py

ilaa-chenjeri-15

Initial Commit

56301ea 6 months ago

Raw History Blame Contribute Delete

392 Bytes

from modules.next_token import get_next_token_probs


def get_next_token(text, model, tokenizer, temperature, rep_penalty):
    tokens, probs, entropy = get_next_token_probs(
        text, model, tokenizer, temperature, top_k=4, rep_penalty=rep_penalty
    )
    high_prob = [tokens[0],probs[0]]
    alternatives = list(zip(tokens[1:], probs[1:]))
    return high_prob, alternatives, entropy

	from modules.next_token import get_next_token_probs


	def get_next_token(text, model, tokenizer, temperature, rep_penalty):
	tokens, probs, entropy = get_next_token_probs(
	text, model, tokenizer, temperature, top_k=4, rep_penalty=rep_penalty
	)
	high_prob = [tokens[0],probs[0]]
	alternatives = list(zip(tokens[1:], probs[1:]))
	return high_prob, alternatives, entropy