Update model.py

model.py

@@ -47,67 +47,10 @@ embeddings_df: Optional[pd.DataFrame] = None
 text_faiss: Optional[object] = None
 image_faiss: Optional[object] = None
 
-# def initialize_models() -> bool:
-#     global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
-    
-#     try:
-#         print(f"Initializing models on device: {device}")
-        
-#         # Initialize CLIP model with error handling
-#         try:
-#             clip_model, _, clip_preprocess = open_clip.create_model_and_transforms(
-#                 'hf-hub:Marqo/marqo-fashionCLIP'
-#             )
-#             clip_model = clip_model.to(device)
-#             clip_model.eval()
-#             clip_tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionCLIP')
-#             print("CLIP model initialized successfully")
-#         except Exception as e:
-#             raise RuntimeError(f"Failed to initialize CLIP model: {str(e)}")
-
-#         # Initialize LLM with optimized settings
-#         try:
-#             model_name = "mistralai/Mistral-7B-v0.1"
-#             quantization_config = BitsAndBytesConfig(
-#                 load_in_4bit=True,
-#                 bnb_4bit_compute_dtype=torch.float16,
-#                 bnb_4bit_use_double_quant=True,
-#                 bnb_4bit_quant_type="nf4"
-#             )
-
-#             # Get token from Streamlit secrets
-#             hf_token = st.secrets["HUGGINGFACE_TOKEN"]
-
-#             llm_tokenizer = AutoTokenizer.from_pretrained(
-#                 model_name,
-#                 padding_side="left",
-#                 truncation_side="left",
-#                 token=hf_token  # Add token here
-#             )
-#             llm_tokenizer.pad_token = llm_tokenizer.eos_token
-
-#             llm_model = AutoModelForCausalLM.from_pretrained(
-#                 model_name,
-#                 quantization_config=quantization_config,
-#                 device_map="auto",
-#                 torch_dtype=torch.float16,
-#                 token=hf_token  # Add token here
-#             )
-#             llm_model.eval()
-#             print("LLM initialized successfully")
-#         except Exception as e:
-#             raise RuntimeError(f"Failed to initialize LLM: {str(e)}")
-
-#         return True
-
-#     except Exception as e:
-#         raise RuntimeError(f"Model initialization failed: {str(e)}")
-
 def initialize_models() -> bool:
     global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
     
     try:
-        device = "cpu"  # Force CPU usage for Streamlit Cloud
         print(f"Initializing models on device: {device}")
         
         # Initialize CLIP model with error handling
@@ -122,10 +65,16 @@ def initialize_models() -> bool:
         except Exception as e:
             raise RuntimeError(f"Failed to initialize CLIP model: {str(e)}")
 
-        # Initialize LLM with CPU-compatible settings
+        # Initialize LLM with optimized settings
         try:
             model_name = "mistralai/Mistral-7B-v0.1"
-            
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4"
+            )
+
             # Get token from Streamlit secrets
             hf_token = st.secrets["HUGGINGFACE_TOKEN"]
 
@@ -133,15 +82,16 @@ def initialize_models() -> bool:
                 model_name,
                 padding_side="left",
                 truncation_side="left",
-                token=hf_token
+                token=hf_token  # Add token here
             )
             llm_tokenizer.pad_token = llm_tokenizer.eos_token
 
             llm_model = AutoModelForCausalLM.from_pretrained(
                 model_name,
+                quantization_config=quantization_config,
                 device_map="auto",
-                low_cpu_mem_usage=True,
-                token=hf_token
+                torch_dtype=torch.float16,
+                token=hf_token  # Add token here
             )
             llm_model.eval()
             print("LLM initialized successfully")