01-ai

/

Yi-6B

like 363

The 6B model has run on generation and the eos_token_id is not set.

Here is the corrected version that produces the intended results:

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the model
model = AutoModelForCausalLM.from_pretrained("01-ai/Yi-6B", device_map="auto", torch_dtype="auto", trust_remote_code=True)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained("01-ai/Yi-6B", trust_remote_code=True)

# Encode the input text
inputs = tokenizer("There's a place where time stands still. A place of breath taking wonder, but also", return_tensors="pt")
max_length = 256 

# Generate output with the end-of-sequence token
outputs = model.generate(
    inputs.input_ids.cuda(),
    max_length=max_length,
    eos_token_id=tokenizer.eos_token_id  # Use the EOS token ID which is 2
)

# Decode and print the output
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

eos_token_id=2 if anyone wants the value

Special Tokens: {'bos_token': '<|startoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '', 'pad_token': ''}
Vocabulary Size: 64000
ID 0 : Token
ID 1 : Token <|startoftext|>
ID 2 : Token <|endoftext|>