Views
No views yet
1%%capture
2%pip install -U bitsandbytes
3%pip install -U transformers
4%pip install -U peft
5%pip install -U accelerate
6%pip install -U trl
7%pip install -U datasets1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
3from peft import PeftModel, PeftConfig
4
5base_model = "google/gemma-7b-it"
6adapter_model = "jojo-ai-mst/rolema-7b-it"
7
8# Load base model(Gemma 7B-it)
9bnbConfig = BitsAndBytesConfig(
10 load_in_4bit = True,
11 bnb_4bit_quant_type="nf4",
12 bnb_4bit_compute_dtype=torch.bfloat16,
13)
14
15model = AutoModelForCausalLM.from_pretrained(base_model,quantization_config=bnbConfig,) # device_map="auto" autosplit for cuda
16model = PeftModel.from_pretrained(model, adapter_model)
17tokenizer = AutoTokenizer.from_pretrained(base_model)
18
19model = model.to("cuda")
20
21inputs = tokenizer("How to learn programming", return_tensors="pt")
22
23inputs = inputs.to("cuda")
24
25outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=1000)
26print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0])