Update app.py

app.py

@@ -1,9 +1,13 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("umar141/Gemma_1B_Baro_v2_vllm")
-model = AutoModelForCausalLM.from_pretrained("umar141/Gemma_1B_Baro_v2_vllm")
+model = AutoModelForCausalLM.from_pretrained(
+    "umar141/Gemma_1B_Baro_v2_vllm",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32  # Use float16 if GPU is available, else float32
+)
 
 # Streamlit page configuration
 st.set_page_config(page_title="Gemma-based Chatbot", page_icon=":robot:")
@@ -27,4 +31,3 @@ if user_input:
     # Display the response
     st.write("AI Response:")
     st.write(response)
-