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.2python import huggingface_hub huggingface_hub.login(<ACCESS_TOKEN>) 1from huggingface_hub import hf_hub_download
2
3model_name = "samvelkoch/masked-fat-mamba" # 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/masked-fat-mamba" # 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 trust_remote_code=True,
11)
12model = AutoModelForCausalLM.from_pretrained(
13 model_name,
14 torch_dtype="auto",
15 device_map={"": "cuda:0"},
16 trust_remote_code=True,
17).cuda().eval()
18
19head_weights = torch.load("classification_head.pth", map_location="cuda")
20# settings can be arbitrary here as we overwrite with saved weights
21head = torch.nn.Linear(1, 1, bias=False).to("cuda")
22head.weight.data = head_weights
23
24inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("cuda")
25
26out = model(**inputs).logits
27
28logits = head(out[:,-1])
29
30print(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): LlamaSdpaAttention(
(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): SiLU()
)
(input_layernorm): LlamaRMSNorm()
(post_attention_layernorm): LlamaRMSNorm()
)
)
(norm): LlamaRMSNorm()
)
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)