tokenizer = AutoTokenizer.from_pretrained(
    model_name, trust_remote_code=True, use_fast=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_name, vocab_size=len(tokenizer), attn_implementation="flash_attention_2",
    trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto"
)
model.eval()

generation_config = GenerationConfig.from_pretrained(
    model_name, do_sample=False, repetition_penalty=1.5,
    min_new_tokens=1, max_new_tokens=128
)

test_input = tokenizer("Wrocław to polski miasto. Wrocław jest ", return_tensors='pt').to(torch.device('cuda'))
test_output = model.generate(**test_input, generation_config=generation_config)
test_preds = tokenizer.batch_decode(sequences=test_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
print(test_preds)