Update app.py

app.py

@@ -11,6 +11,9 @@ from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from dotenv import load_dotenv
 from htmlTemplates import css
 
+# Set Streamlit page configuration
+st.set_page_config(page_title="Chat with Notes and AI", page_icon=":books:", layout="wide")
+
 # Load environment variables
 load_dotenv()
 
@@ -23,7 +26,7 @@ def load_pipeline():
     tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left", trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         model_name, 
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Use float16 for GPU, float32 for CPU
+        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
     )
@@ -33,7 +36,7 @@ def load_pipeline():
         task="text-generation", 
         model=model, 
         tokenizer=tokenizer, 
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         device_map="auto", 
         return_full_text=True  # Required for LangChain compatibility
     )
@@ -100,7 +103,6 @@ def handle_question(question, vectorstore=None):
     return llm_chain.run(instruction=question)
 
 def main():
-    st.set_page_config(page_title="Chat with Notes and AI", page_icon=":books:", layout="wide")
     st.write(css, unsafe_allow_html=True)
 
     # Initialize session state