Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import transformers
3import torch
4
5model_id = "aeonium/Aeonium-v1.1-Chat-4B"
6
7tokenizer = AutoTokenizer.from_pretrained(model_id)
8model = AutoModelForCausalLM.from_pretrained(
9 model_id,
10 device_map="cuda",
11 torch_dtype=torch.bfloat16,
12)
13
14chat = [
15 { "role": "user", "content": "Привет!" },
16]
17
18prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
19inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
20outputs = model.generate(input_ids=inputs, max_new_tokens=48)
21print(tokenizer.decode(outputs[0]))