Update app.py

app.py

@@ -7,6 +7,8 @@ from groq import Groq
 from PyPDF2 import PdfReader
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 # Set up environment variables
 os.environ["GROQ_API_KEY"] = "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"
@@ -18,14 +20,37 @@ client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 whisper_model = whisper.load_model("base")  # You can choose other models like "small", "medium", "large"
 
 # Initialize the tokenizer and model from the saved checkpoint for RAG
+# Updated model loading code with disk offloading
+
+
+# Specify the folder where offloaded model parts will be stored
+offload_folder = "./offload"
+
+# Ensure the offload folder exists
+os.makedirs(offload_folder, exist_ok=True)
+
+# Initialize the tokenizer
 rag_tokenizer = AutoTokenizer.from_pretrained("himmeow/vi-gemma-2b-RAG")
-rag_model = AutoModelForCausalLM.from_pretrained(
+
+# Initialize empty weights with the specified offload folder
+with init_empty_weights():
+    rag_model = AutoModelForCausalLM.from_pretrained(
+        "himmeow/vi-gemma-2b-RAG",
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        offload_folder=offload_folder
+    )
+
+# Dispatch the model with offloading options
+rag_model = load_checkpoint_and_dispatch(
+    rag_model,
     "himmeow/vi-gemma-2b-RAG",
     device_map="auto",
-    torch_dtype=torch.bfloat16
+    offload_folder=offload_folder,
+    offload_state_dict=True
 )
 
-# Use GPU if available for RAG model
+# Use GPU if available
 if torch.cuda.is_available():
     rag_model.to("cuda")