Update app.py

app.py

@@ -1,27 +1,16 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, disk_offload
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Load the Phi 2 model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(
-    "microsoft/phi-2",
-    trust_remote_code=True
-)
-
-model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/phi-2",
-    device_map="auto",
-    trust_remote_code=True
-)
-
-# Offload the model to disk
-model, tokenizer = disk_offload(model, tokenizer)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", device_map="auto", trust_remote_code=True)
 
 # Streamlit UI
 st.title("Microsoft Phi 2 Streamlit App")
 
 # User input prompt
-prompt = st.text_area("Enter your prompt:", """Write a story about Nasa""")
+prompt = st.text_area("Enter your prompt:", "Write a story about Nasa")
 
 # Generate output based on user input
 if st.button("Generate Output"):