Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3import os
4
5device = "cuda"
6device_map = {
7 "lm_head": device,
8 "model": device
9}
10
11model_name = "piotr-ai/polanka-qwen2-3b-v0.3-chat"
12
13tokenizer = AutoTokenizer.from_pretrained(model_name)
14model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map=device_map)
15
16prompt = "Co to jest psychologia?"
17messages = [
18 {"role": "system", "content": "Jesteś przydatnym asystentem."},
19 {"role": "assistant", "content": "Witam, jak mogę pomóc?"},
20 {"role": "user", "content": prompt}
21]
22text = tokenizer.apply_chat_template(messages, tokenize=False)
23model_inputs = tokenizer([text], return_tensors="pt").to(device)
24
25generated_ids = model.generate(
26 model_inputs.input_ids,
27 max_new_tokens=256,
28 num_return_sequences=1,
29 eos_token_id=tokenizer.convert_tokens_to_ids('<|im_end|>'),
30 pad_token_id=tokenizer.eos_token_id
31)
32
33print(tokenizer.decode(generated_ids[0]))<|im_start|>system
Jesteś przydatnym asystentem.<|im_end|>
<|im_start|>assistant
Witam, jak mogę pomóc?<|im_end|>
<|im_start|>user
Co to jest psychologia?<|im_end|>
<|im_start|>assistant
Psychologia to nauka badająca zachowanie i procesy umysłowe. Jest to gałąź psychologii, która koncentruje się na zrozumieniu i wyjaśnianiu zachowań i procesów umysłowych.<|im_end|>