Spaces:
Build error
Build error
Wisdom Chen
commited on
Update model.py
Browse files
model.py
CHANGED
@@ -47,12 +47,6 @@ text_faiss: Optional[object] = None
|
|
47 |
image_faiss: Optional[object] = None
|
48 |
|
49 |
def initialize_models() -> bool:
|
50 |
-
"""
|
51 |
-
Initialize CLIP and LLM models with proper error handling and GPU optimization.
|
52 |
-
|
53 |
-
Returns:
|
54 |
-
bool: True if initialization successful, raises RuntimeError otherwise
|
55 |
-
"""
|
56 |
global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
|
57 |
|
58 |
try:
|
@@ -72,21 +66,26 @@ def initialize_models() -> bool:
|
|
72 |
|
73 |
# Initialize LLM with optimized settings
|
74 |
try:
|
75 |
-
|
76 |
-
|
|
|
|
|
|
|
|
|
|
|
77 |
model_name = "mistralai/Mistral-7B-v0.1"
|
78 |
quantization_config = BitsAndBytesConfig(
|
79 |
load_in_4bit=True,
|
80 |
bnb_4bit_compute_dtype=torch.float16,
|
81 |
bnb_4bit_use_double_quant=True,
|
82 |
-
bnb_4bit_quant_type="nf4"
|
83 |
-
token=os.environ.get("HF_TOKEN")
|
84 |
)
|
85 |
|
86 |
llm_tokenizer = AutoTokenizer.from_pretrained(
|
87 |
model_name,
|
88 |
padding_side="left",
|
89 |
-
truncation_side="left"
|
|
|
90 |
)
|
91 |
llm_tokenizer.pad_token = llm_tokenizer.eos_token
|
92 |
|
@@ -94,7 +93,8 @@ def initialize_models() -> bool:
|
|
94 |
model_name,
|
95 |
quantization_config=quantization_config,
|
96 |
device_map="auto",
|
97 |
-
torch_dtype=torch.float16
|
|
|
98 |
)
|
99 |
llm_model.eval()
|
100 |
print("LLM initialized successfully")
|
|
|
47 |
image_faiss: Optional[object] = None
|
48 |
|
49 |
def initialize_models() -> bool:
|
|
|
|
|
|
|
|
|
|
|
|
|
50 |
global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
|
51 |
|
52 |
try:
|
|
|
66 |
|
67 |
# Initialize LLM with optimized settings
|
68 |
try:
|
69 |
+
# Check for HF_TOKEN and authenticate
|
70 |
+
hf_token = os.environ.get("HF_TOKEN")
|
71 |
+
if not hf_token:
|
72 |
+
raise RuntimeError("HF_TOKEN environment variable is not set")
|
73 |
+
|
74 |
+
login(token=hf_token)
|
75 |
+
|
76 |
model_name = "mistralai/Mistral-7B-v0.1"
|
77 |
quantization_config = BitsAndBytesConfig(
|
78 |
load_in_4bit=True,
|
79 |
bnb_4bit_compute_dtype=torch.float16,
|
80 |
bnb_4bit_use_double_quant=True,
|
81 |
+
bnb_4bit_quant_type="nf4"
|
|
|
82 |
)
|
83 |
|
84 |
llm_tokenizer = AutoTokenizer.from_pretrained(
|
85 |
model_name,
|
86 |
padding_side="left",
|
87 |
+
truncation_side="left",
|
88 |
+
token=hf_token
|
89 |
)
|
90 |
llm_tokenizer.pad_token = llm_tokenizer.eos_token
|
91 |
|
|
|
93 |
model_name,
|
94 |
quantization_config=quantization_config,
|
95 |
device_map="auto",
|
96 |
+
torch_dtype=torch.float16,
|
97 |
+
token=hf_token
|
98 |
)
|
99 |
llm_model.eval()
|
100 |
print("LLM initialized successfully")
|