Views
No views yet
transformers library on a machine with GPUs, first make sure you have the transformers library installed.pip install transformers==4.40.1token=True in the pipeline and login to hugginface_hub by running1import huggingface_hub
2huggingface_hub.login(<ACCESS_TOKEN>)token in the pipeline1from transformers import pipeline
2
3generate_text = pipeline(
4 model="yifanxie/malachite-gibbon-step2",
5 torch_dtype="auto",
6 trust_remote_code=True,
7 use_fast=True,
8 device_map={"": "cuda:0"},
9 token=True,
10)
11
12# generate configuration can be modified to your needs
13# generate_text.model.generation_config.min_new_tokens = 2
14# generate_text.model.generation_config.max_new_tokens = 256
15# generate_text.model.generation_config.do_sample = False
16# generate_text.model.generation_config.num_beams = 1
17# generate_text.model.generation_config.temperature = float(0.0)
18# generate_text.model.generation_config.repetition_penalty = float(1.0)
19
20messages = [
21 {
22 "role": "system",
23 "content": "You are a friendly and polite chatbot.",
24 },
25 {"role": "user", "content": "Hi, how are you?"},
26 {"role": "assistant", "content": "I'm doing great, how about you?"},
27 {"role": "user", "content": "Why is drinking water so healthy?"},
28]
29
30res = generate_text(
31 messages,
32 renormalize_logits=True
33)
34print(res[0]["generated_text"][-1]['content'])1print(generate_text.tokenizer.apply_chat_template(
2 messages,
3 tokenize=False,
4 add_generation_prompt=True,
5))1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "yifanxie/malachite-gibbon-step2" # either local folder or huggingface model name
4# Important: The prompt needs to be in the same format the model was trained with.
5# You can find an example prompt in the experiment logs.
6messages = [
7 {
8 "role": "system",
9 "content": "You are a friendly and polite chatbot.",
10 },
11 {"role": "user", "content": "Hi, how are you?"},
12 {"role": "assistant", "content": "I'm doing great, how about you?"},
13 {"role": "user", "content": "Why is drinking water so healthy?"},
14]
15
16tokenizer = AutoTokenizer.from_pretrained(
17 model_name,
18 use_fast=True,
19 trust_remote_code=True,
20)
21model = AutoModelForCausalLM.from_pretrained(
22 model_name,
23 torch_dtype="auto",
24 device_map={"": "cuda:0"},
25 trust_remote_code=True,
26)
27model.cuda().eval()
28
29# generate configuration can be modified to your needs
30# model.generation_config.min_new_tokens = 2
31# model.generation_config.max_new_tokens = 256
32# model.generation_config.do_sample = False
33# model.generation_config.num_beams = 1
34# model.generation_config.temperature = float(0.0)
35# model.generation_config.repetition_penalty = float(1.0)
36
37inputs = tokenizer.apply_chat_template(
38 messages,
39 tokenize=True,
40 add_generation_prompt=True,
41 return_tensors="pt",
42 return_dict=True,
43).to("cuda")
44
45tokens = model.generate(
46 input_ids=inputs["input_ids"],
47 attention_mask=inputs["attention_mask"],
48 renormalize_logits=True
49)[0]
50
51tokens = tokens[inputs["input_ids"].shape[1]:]
52answer = tokenizer.decode(tokens, skip_special_tokens=True)
53print(answer)load_in_8bit=True or load_in_4bit=True. Also, sharding on multiple GPUs is possible by setting device_map=auto.GemmaForCausalLM(
(model): GemmaModel(
(embed_tokens): Embedding(256000, 2048, padding_idx=0)
(layers): ModuleList(
(0-17): 18 x GemmaDecoderLayer(
(self_attn): GemmaSdpaAttention(
(q_proj): Linear(in_features=2048, out_features=2048, bias=False)
(k_proj): Linear(in_features=2048, out_features=256, bias=False)
(v_proj): Linear(in_features=2048, out_features=256, bias=False)
(o_proj): Linear(in_features=2048, out_features=2048, bias=False)
(rotary_emb): GemmaRotaryEmbedding()
)
(mlp): GemmaMLP(
(gate_proj): Linear(in_features=2048, out_features=16384, bias=False)
(up_proj): Linear(in_features=2048, out_features=16384, bias=False)
(down_proj): Linear(in_features=16384, out_features=2048, bias=False)
(act_fn): PytorchGELUTanh()
)
(input_layernorm): GemmaRMSNorm()
(post_attention_layernorm): GemmaRMSNorm()
)
)
(norm): GemmaRMSNorm()
)
(lm_head): Linear(in_features=2048, out_features=256000, bias=False)
)