mirror of
https://git.victorphan.net/basketballcantho/book_translator.git
synced 2026-08-05 06:33:11 +07:00
fix bug
This commit is contained in:
+3
-2
@@ -205,7 +205,7 @@ def load_model(model_id: str, load_in_8bit: bool = True) -> None:
|
||||
global _model, _tokenizer
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
|
||||
print(f"\n🔄 Loading model: {model_id}")
|
||||
print(f" Quantization : {'INT8 (bitsandbytes)' if load_in_8bit else 'auto (float16/bfloat16)'}")
|
||||
@@ -221,10 +221,11 @@ def load_model(model_id: str, load_in_8bit: bool = True) -> None:
|
||||
|
||||
if load_in_8bit:
|
||||
# INT8 quantization via bitsandbytes — fits ~72 GB for a 72B model
|
||||
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_id,
|
||||
device_map="auto",
|
||||
load_in_8bit=True,
|
||||
quantization_config=quantization_config,
|
||||
trust_remote_code=True,
|
||||
torch_dtype=torch.float16,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user