Update app.py

app.py

@@ -2,32 +2,29 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
+# Use GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Move model to CPU
-model.to("cpu")
+st.title("Text Generation with Hugging Face Transformers")
 
-# Streamlit app
-st.title("Text Generation with Transformers")
+# Input prompt from user
+prompt = st.text_area("Enter a prompt:", "this news is real pyresearch given right computer vision videos?")
 
-# Input prompt
-prompt = st.text_input("Enter your prompt:")
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", device=device, trust_remote_code=True)
 
-# Generate button
+# Generate text on button click
 if st.button("Generate"):
     with torch.no_grad():
-        # Tokenize and generate output
-        token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+        token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(device)
         output_ids = model.generate(
-            token_ids.to(model.device),
+            token_ids,
             max_new_tokens=512,
             do_sample=True,
             temperature=0.1
         )
-    
-    # Decode and display the generated text
+
     generated_text = tokenizer.decode(output_ids[0][token_ids.size(1):])
     st.text("Generated Text:")
-    st.text(generated_text)
+    st.write(generated_text)