Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
2import torch
3
4TORCH_DTYPE = 'bfloat16'
5nf4_config = BitsAndBytesConfig(
6 load_in_4bit=True,
7 bnb_4bit_quant_type='nf4',
8 bnb_4bit_use_double_quant=True,
9 bnb_4bit_compute_dtype=getattr(torch, TORCH_DTYPE)
10)
11
12tokenizer = AutoTokenizer.from_pretrained('mesolitica/mallam-5B-4096')
13model = AutoModelForCausalLM.from_pretrained(
14 'mesolitica/mallam-5B-4096',
15 use_flash_attention_2 = True,
16 quantization_config = nf4_config
17)
18prompt = '<s>nama saya'
19inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')
20
21generate_kwargs = dict(
22 inputs,
23 max_new_tokens=512,
24 top_p=0.95,
25 top_k=50,
26 temperature=0.9,
27 do_sample=True,
28 num_beams=1,
29 repetition_penalty=1.05,
30)
31r = model.generate(**generate_kwargs)