Uploading Mulitimodal Retrieval Augmented Generation System.

model_setup.py

@@ -5,6 +5,7 @@ import gc
 
 from sentence_transformers import SentenceTransformer
 from transformers import AutoProcessor, Gemma3ForConditionalGeneration, BitsAndBytesConfig
+from accelerate import disk_offload
 from utils import clear_gpu_cache
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -12,25 +13,26 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Embedding model
 embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-# Gemma3 quantized config
+# Gemma3 quantization config
 model_name = "google/gemma-3-4b-it"
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.bfloat16,
     bnb_4bit_use_double_quant=True,
     bnb_4bit_quant_type="nf4",
-    llm_int8_enable_fp32_cpu_offload=True  # Allow offloading    
+    # llm_int8_enable_fp32_cpu_offload=True  # Allow offloading    
 )
 
 # Load Gemma3
 model = Gemma3ForConditionalGeneration.from_pretrained(
     model_name,
     torch_dtype=torch.bfloat16,
-    device_map="auto",
+    device_map="cpu", # not "auto" since there is no GPU
     quantization_config=bnb_config,
     low_cpu_mem_usage=True,
-    attn_implementation="sdpa"
+    # attn_implementation="sdpa"
 )
+disk_offload(model=model, offload_dir="offload")
 model.eval()
 
 # Processor