Updated app.py

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# torch.set_default_device("cuda")
 # Load the Phi 2 model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
     "microsoft/phi-2",
@@ -10,8 +9,8 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 
 model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/phi-2",
-    device_map="cpu",
+    "kroonen/phi-2-GGUF",
+    # device_map="cpu",
     trust_remote_code=True,
     torch_dtype=torch.float32
 )
@@ -38,4 +37,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)