anhvv200053 commited on
Commit
dc745aa
·
verified ·
1 Parent(s): 92406c0

Update model_setup.py

Browse files
Files changed (1) hide show
  1. model_setup.py +25 -25
model_setup.py CHANGED
@@ -1,25 +1,25 @@
1
- import torch
2
- from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
3
-
4
- # Cấu hình BitsAndBytes để tải mô hình 4-bit
5
- bnb_config = BitsAndBytesConfig(
6
- load_in_4bit=True,
7
- bnb_4bit_quant_type='nf4',
8
- bnb_4bit_compute_dtype="float16",
9
- bnb_4bit_use_double_quant=False,
10
- )
11
-
12
- # Thiết lập mô hình và tokenizer
13
- def load_model():
14
- model = AutoModelForCausalLM.from_pretrained(
15
- "anhvv200053/Vinallama-2-7B-updated1-instruction-v2",
16
- quantization_config=bnb_config,
17
- device_map={"": 0},
18
- use_auth_token=True
19
- )
20
- model.config.pretraining_tp = 1
21
-
22
- tokenizer = AutoTokenizer.from_pretrained('anhvv200053/Vinallama-2-7B-updated1-instruction-v2', trust_remote_code=True, use_fast=True)
23
- tokenizer.pad_token = tokenizer.eos_token
24
-
25
- return model, tokenizer
 
1
+ import torch
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
3
+
4
+ # Cấu hình BitsAndBytes để tải mô hình 4-bit
5
+ bnb_config = BitsAndBytesConfig(
6
+ load_in_4bit=True,
7
+ bnb_4bit_quant_type='nf4',
8
+ bnb_4bit_compute_dtype="float16",
9
+ bnb_4bit_use_double_quant=False,
10
+ )
11
+
12
+ # Thiết lập mô hình và tokenizer
13
+ def load_model():
14
+ model = AutoModelForCausalLM.from_pretrained(
15
+ "anhvv200053/Vinallama-2-7B-updated1-instruction-v2",
16
+ quantization_config=bnb_config,
17
+ device_map={"": 0},
18
+ token = token
19
+ )
20
+ model.config.pretraining_tp = 1
21
+
22
+ tokenizer = AutoTokenizer.from_pretrained('anhvv200053/Vinallama-2-7B-updated1-instruction-v2', trust_remote_code=True, use_fast=True, token = token)
23
+ tokenizer.pad_token = tokenizer.eos_token
24
+
25
+ return model, tokenizer