Update app.py

app.py

@@ -10,6 +10,13 @@ import requests
 import time
 from tqdm import tqdm
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from vllm import LLM, SamplingParams
+
+@st.cache_resource
+def load_model():
+    return LLM(model="deepseek-ai/DeepSeek-R1")
+
+llm = load_model()
 
 def sorting(df):
     df.index = list(map(float, df.index))
@@ -1017,11 +1024,15 @@ try:
     elif main_option == "AI Chat":
         prompt_user = st.text_input("Write a prompt.")
 
-        messages = [
-            {"role": "user", "content": prompt_user},
-        ]
-        pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)
-        st.markdown(pipe(messages))
+        if st.button("Send"):
+            if user_input:
+                sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
+                output = llm.generate([prompt_user], sampling_params)
+                response = output[0].outputs[0].text
+                
+                st.text_area("Answer:", response, height=200)
+            else:
+                st.warning("Enter a prompt:")
 
 except Exception as e:
         st.error(f"❌ Error: {e}")