Update app.py

app.py

@@ -1,10 +1,7 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from transformers import AutoTokenizer
-
+from transformers import AutoTokenizer, AutoModelForCausalLM, disk_offload
 import torch
 
-
 # Load the Phi 2 model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
     "microsoft/phi-2",
@@ -17,6 +14,9 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 
+# Offload the model to disk
+model, tokenizer = disk_offload(model, tokenizer)
+
 # Streamlit UI
 st.title("Microsoft Phi 2 Streamlit App")
 
@@ -36,4 +36,4 @@ if st.button("Generate Output"):
 
     output = tokenizer.decode(output_ids[0][token_ids.size(1):])
     st.text("Generated Output:")
-    st.write(output)
+    st.write(output)