Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GemmaTokenizer
3
4# download the file or use the HF_Token to get the model
5model_id = file
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_quant_type="nf4",
9 bnb_4bit_compute_dtype=torch.bfloat16
10)
11
12tokenizer = AutoTokenizer.from_pretrained(model_id)
13model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0})
14
15text = "sociological imagination is "
16device = "cuda:0"
17inputs = tokenizer(text, return_tensors="pt").to(device)
18
19outputs = model.generate(**inputs, max_new_tokens=60)
20print(tokenizer.decode(outputs[0], skip_special_tokens=True))
21
22# Output: sociological imagination is <strong>the ability to see the relationship between personal troubles and public issues
23