Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2import torch
3
4model_id = "isemmanuelolowe/Jamba-8xMoE_slerp"
5
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7quantization_config = BitsAndBytesConfig(
8 load_in_4bit=True,
9 # load_in_8bit=True,
10 bnb_4bit_quant_type="nf4",
11 bnb_4bit_compute_dtype=torch.bfloat16,
12 bnb_4bit_use_double_quant=True,
13 llm_int8_skip_modules=["mamba"],
14)
15
16model = AutoModelForCausalLM.from_pretrained(
17 model_id,
18 trust_remote_code=True,
19 torch_dtype=torch.bfloat16,
20 attn_implementation="flash_attention_2",
21 quantization_config=quantization_config
22)
23
24input_ids = tokenizer("Here is how to do bubble sort\n```python\n", return_tensors="pt")["input_ids"].to("cuda")
25
26out = model.generate(input_ids, max_new_tokens=256, temperature=0, repetition_penalty=1)
27print(tokenizer.batch_decode(out, skip_special_tokens=True))['Here is how to do bubble sort\n```python\ndef bubble_sort(array):\n for i in 0, len(array):\n for j in 0, len(array):\n if a[i] < a[j]\n a[i], a[j]\n\n```\n\n\n\n\n\n\n']