Update app.py

app.py

@@ -19,25 +19,28 @@ load_dotenv()
 # Dolly-v2-3b model pipeline
 @st.cache_resource
 def load_pipeline():
-    # Use recommended settings for Dolly-v2-3b
     model_name = "databricks/dolly-v2-3b"
 
+    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left", trust_remote_code=True)
+
+    # Load model with offload folder for disk storage of weights
     model = AutoModelForCausalLM.from_pretrained(
         model_name, 
         torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,  # Use float16 for GPU, float32 for CPU
         device_map="auto",          # Automatically map model to available devices (e.g., GPU if available)
-        trust_remote_code=True
+        trust_remote_code=True,
+        offload_folder="./offload_weights"  # Folder to store offloaded weights
     )
 
-    # Load the pipeline with required configurations
+    # Return text-generation pipeline
     return pipeline(
         task="text-generation", 
         model=model, 
         tokenizer=tokenizer, 
         torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         device_map="auto", 
-        return_full_text=True  # Required for LangChain compatibility
+        return_full_text=True
     )
 
 # Initialize Dolly pipeline