Update app.py

app.py

@@ -9,8 +9,7 @@ import plotly.figure_factory as ff
 import requests
 import time
 from tqdm import tqdm
-from transformers import pipeline
-
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 
 def sorting(df):
     df.index = list(map(float, df.index))
@@ -1017,10 +1016,15 @@ try:
 
     elif main_option == "AI Chat":
         prompt_user = st.text_input("Write a prompt.")
+
+        model_name = "deepseek-ai/DeepSeek-R1"
+        model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
         messages = [
             {"role": "user", "content": prompt_user},
         ]
-        pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True, device_map="auto") # Added device_map
         st.markdown(pipe(messages))
 
 except Exception as e: