Views
No views yet
transformers library on a machine with GPUs, first make sure you have the transformers library installed.pip install transformers==4.34.0python import huggingface_hub huggingface_hub.login(<ACCES_TOKEN>) 1from huggingface_hub import hf_hub_download
2
3model_name = "samvelkoch/friendly-mouse" # either local folder or huggingface model name
4hf_hub_download(repo_id=model_name, filename="classification_head.pth", local_dir="./")1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "samvelkoch/friendly-mouse" # 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.
6prompt = "How are you?"
7
8tokenizer = AutoTokenizer.from_pretrained(
9 model_name,
10 use_fast=True,
11 trust_remote_code=True,
12)
13model = AutoModelForCausalLM.from_pretrained(
14 model_name,
15 torch_dtype="auto",
16 device_map={"": "cuda:0"},
17 trust_remote_code=True,
18).cuda().eval()
19
20head_weights = torch.load("classification_head.pth", map_location="cuda")
21# settings can be arbitrary here as we overwrite with saved weights
22head = torch.nn.Linear(1, 1, bias=False).to("cuda")
23head.weight.data = head_weights
24
25inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("cuda")
26
27out = model(**inputs).logits
28
29logits = head(out[:,-1])
30
31print(logits)load_in_8bit=True or load_in_4bit=True. Also, sharding on multiple GPUs is possible by setting device_map=auto.LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(32000, 4096, padding_idx=0)
(layers): ModuleList(
(0-31): 32 x LlamaDecoderLayer(
(self_attn): LlamaAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(rotary_emb): LlamaRotaryEmbedding()
)
(mlp): LlamaMLP(
(gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
(up_proj): Linear(in_features=4096, out_features=11008, bias=False)
(down_proj): Linear(in_features=11008, out_features=4096, bias=False)
(act_fn): SiLUActivation()
)
(input_layernorm): LlamaRMSNorm()
(post_attention_layernorm): LlamaRMSNorm()
)
)
(norm): LlamaRMSNorm()
)
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)