message_edits

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+.vscode

app.py

@@ -1,38 +1,50 @@
-import streamlit as st 
+import streamlit as st
 from transformers import pipeline
 
-## Model stuff
-model_name = "deepseek-ai/DeepSeek-R1"
+# Model stuff
+model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 messages = [
-    {"role": "System", "content": "You are a AI model that hits on people when they talk to you."},
+    {"role": "System", "content": "You are an AI model that responds in a flirtatious and playful manner."},
 ]
 
-@st.cache_data
-def load_model():
-    pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")
-    return pipeline
+@st.cache_resource
+def load_model(model_name):
+    return pipeline("text-generation", model=model_name, device_map="auto")
 
-def response(content: str):
+pipe = load_model(model_name)
+
+def generate_response(prompt: str):
+    messages.append({"role": "user", "content": prompt})
     try:
-        pipe = load_model()
-        messages.append({"role": "user", "content": str})
-        content = pipe(messages, temperature=0.7, top_p=0.95)
-        response = content[0]['generated_text']
-        messages.append({"role": "assistant", "answer_to_previous_question": response_text})
-        return response
+        content = pipe(messages, max_new_tokens=200, temperature=0.7, top_p=0.95)
+        response_text = content[0]['generated_text']
+        messages.append({"role": "assistant", "content": response_text})
+        return response_text
     except Exception as e:
         st.error(f"An error occurred: {e}")
         return "I'm having a little trouble right now. Please try again later."
 
 
-## Streamlit Start :)
-st.title('Deepseek AI Partner')
+# Streamlit Start :)
+st.title('DeepSeek AI Partner')
+
+if "messages" not in st.session_state:
+    st.session_state["messages"] = messages
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
 if prompt := st.chat_input("What is up?"):
-    # Display user message in chat message container
     with st.chat_message("user"):
         st.markdown(prompt)
-    if prompt:
-        with st.chat_message("assistant"):
-            respone = response(prompt)
-            st.markdown(response)
+
+    with st.chat_message("assistant"):
+        response = generate_response(prompt)
+        st.markdown(response)
+
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.session_state.messages.append({"role": "assistant", "content": response})
+
+if __name__ == "__main__":
+    st.run()