Spaces:
Build error
Build error
Wisdom Chen
commited on
Update model.py
Browse files
model.py
CHANGED
@@ -51,17 +51,31 @@ def initialize_models() -> bool:
|
|
51 |
global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
|
52 |
|
53 |
try:
|
54 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
55 |
|
56 |
# Initialize LLM with optimized settings
|
57 |
try:
|
58 |
model_name = "mistralai/Mistral-7B-v0.1"
|
59 |
quantization_config = BitsAndBytesConfig(
|
60 |
load_in_4bit=True,
|
61 |
-
|
62 |
-
bnb_4bit_use_double_quant=True
|
|
|
63 |
)
|
64 |
|
|
|
65 |
hf_token = st.secrets.get("HUGGINGFACE_TOKEN")
|
66 |
if not hf_token:
|
67 |
raise ValueError("HUGGINGFACE_TOKEN not found in Streamlit secrets")
|
@@ -70,19 +84,17 @@ def initialize_models() -> bool:
|
|
70 |
llm_tokenizer = AutoTokenizer.from_pretrained(
|
71 |
model_name,
|
72 |
token=hf_token,
|
73 |
-
trust_remote_code=True
|
74 |
-
revision="v0.1"
|
75 |
)
|
76 |
llm_tokenizer.pad_token = llm_tokenizer.eos_token
|
77 |
-
|
78 |
llm_model = AutoModelForCausalLM.from_pretrained(
|
79 |
model_name,
|
80 |
token=hf_token,
|
81 |
quantization_config=quantization_config,
|
82 |
device_map="auto",
|
83 |
torch_dtype=torch.float16,
|
84 |
-
trust_remote_code=True
|
85 |
-
revision="v0.1"
|
86 |
)
|
87 |
llm_model.eval()
|
88 |
print("LLM initialized successfully")
|
|
|
51 |
global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
|
52 |
|
53 |
try:
|
54 |
+
print(f"Initializing models on device: {device}")
|
55 |
+
|
56 |
+
# Initialize CLIP model with error handling
|
57 |
+
try:
|
58 |
+
clip_model, _, clip_preprocess = open_clip.create_model_and_transforms(
|
59 |
+
'hf-hub:Marqo/marqo-fashionCLIP'
|
60 |
+
)
|
61 |
+
clip_model = clip_model.to(device)
|
62 |
+
clip_model.eval()
|
63 |
+
clip_tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionCLIP')
|
64 |
+
print("CLIP model initialized successfully")
|
65 |
+
except Exception as e:
|
66 |
+
raise RuntimeError(f"Failed to initialize CLIP model: {str(e)}")
|
67 |
|
68 |
# Initialize LLM with optimized settings
|
69 |
try:
|
70 |
model_name = "mistralai/Mistral-7B-v0.1"
|
71 |
quantization_config = BitsAndBytesConfig(
|
72 |
load_in_4bit=True,
|
73 |
+
bnb_4bit_compute_dtype=torch.float16,
|
74 |
+
bnb_4bit_use_double_quant=True,
|
75 |
+
bnb_4bit_quant_type="nf4"
|
76 |
)
|
77 |
|
78 |
+
# Get token from Streamlit secrets
|
79 |
hf_token = st.secrets.get("HUGGINGFACE_TOKEN")
|
80 |
if not hf_token:
|
81 |
raise ValueError("HUGGINGFACE_TOKEN not found in Streamlit secrets")
|
|
|
84 |
llm_tokenizer = AutoTokenizer.from_pretrained(
|
85 |
model_name,
|
86 |
token=hf_token,
|
87 |
+
trust_remote_code=True
|
|
|
88 |
)
|
89 |
llm_tokenizer.pad_token = llm_tokenizer.eos_token
|
90 |
+
|
91 |
llm_model = AutoModelForCausalLM.from_pretrained(
|
92 |
model_name,
|
93 |
token=hf_token,
|
94 |
quantization_config=quantization_config,
|
95 |
device_map="auto",
|
96 |
torch_dtype=torch.float16,
|
97 |
+
trust_remote_code=True
|
|
|
98 |
)
|
99 |
llm_model.eval()
|
100 |
print("LLM initialized successfully")
|