Views
No views yet
transformers library on a machine with GPUs, first make sure you have the transformers library installed.pip install transformers==4.36.1python import huggingface_hub huggingface_hub.login(<ACCESS_TOKEN>) 1from huggingface_hub import hf_hub_download
2
3model_name = "aghorbani/bank-tx-cat-opt-125m" # 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 = "aghorbani/bank-tx-cat-opt-125m" # 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.OPTForCausalLM(
(model): OPTModel(
(decoder): OPTDecoder(
(embed_tokens): Embedding(50272, 768, padding_idx=1)
(embed_positions): OPTLearnedPositionalEmbedding(2050, 768)
(final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
(layers): ModuleList(
(0-11): 12 x OPTDecoderLayer(
(self_attn): OPTAttention(
(k_proj): Linear(in_features=768, out_features=768, bias=True)
(v_proj): Linear(in_features=768, out_features=768, bias=True)
(q_proj): Linear(in_features=768, out_features=768, bias=True)
(out_proj): Linear(in_features=768, out_features=768, bias=True)
)
(activation_fn): ReLU()
(self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=768, out_features=3072, bias=True)
(fc2): Linear(in_features=3072, out_features=768, bias=True)
(final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
)
)
)
)
(lm_head): Linear(in_features=768, out_features=50272, bias=False)
)