Update my_model/KBVQA.py
Browse files- my_model/KBVQA.py +12 -2
my_model/KBVQA.py
CHANGED
|
@@ -68,8 +68,18 @@ class KBVQA():
|
|
| 68 |
|
| 69 |
def load_fine_tuned_model(self):
|
| 70 |
|
| 71 |
-
self.kbvqa_model = AutoModelForCausalLM.from_pretrained(self.kbvqa_model_name,
|
| 72 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 73 |
|
| 74 |
|
| 75 |
@property
|
|
|
|
| 68 |
|
| 69 |
def load_fine_tuned_model(self):
|
| 70 |
|
| 71 |
+
self.kbvqa_model = AutoModelForCausalLM.from_pretrained(self.kbvqa_model_name,
|
| 72 |
+
device_map="auto",
|
| 73 |
+
low_cpu_mem_usage=True,
|
| 74 |
+
quantization_config=self.bnb_config,
|
| 75 |
+
token=self.access_token)
|
| 76 |
+
|
| 77 |
+
self.kbvqa_tokenizer = AutoTokenizer.from_pretrained(self.kbvqa_model_name,
|
| 78 |
+
use_fast=self.use_fast,
|
| 79 |
+
low_cpu_mem_usage=True,
|
| 80 |
+
trust_remote_code=self.trust_remote,
|
| 81 |
+
add_eos_token=self.add_eos_token,
|
| 82 |
+
token=self.access_token)
|
| 83 |
|
| 84 |
|
| 85 |
@property
|