Update app.py

app.py

@@ -1,5 +1,9 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import accelerate
+
+accelerator = accelerate.Accelerator(device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map=accelerator.device_map)
 
 # Load the model and tokenizer
 @st.cache_resource
@@ -31,6 +35,7 @@ def main():
         with st.spinner("Generating response..."):
             response = generate_response(prompt)
             st.write(response)
+            
 
 if __name__ == "__main__":
     main()