Update app.py

app.py

@@ -21,3 +21,74 @@ else:
     print("CUDA is not available. Using CPU.")
 
 print(f"Using device: {device}")
+
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+# Function to generate chatbot response
+def generate_response(prompt, chat_history_ids=None):
+    inputs = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors="pt")
+
+    if chat_history_ids is None:
+        chat_history_ids = None
+    else:
+        chat_history_ids = torch.tensor(chat_history_ids)
+
+    # generate a response while limiting the total chat history to 1000 tokens,
+    chat_history_ids = model.generate(
+        inputs, max_length=1000,
+        pad_token_id=tokenizer.eos_token_id,
+        no_repeat_ngram_size=3,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95,
+        chat_history_ids = chat_history_ids
+    )
+
+    response = tokenizer.decode(chat_history_ids[:, inputs.shape[-1]:][0], skip_special_tokens=True)
+    return response, chat_history_ids.tolist()
+
+# Streamlit app
+st.title("Simple Chatbot")
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "chat_history_ids" not in st.session_state:
+    st.session_state.chat_history_ids = None
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Accept user input
+if prompt := st.chat_input("What is up?"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    # Display user message in chat message container
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    # Generate and display chatbot response
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        full_response = ""
+        response, st.session_state.chat_history_ids = generate_response(prompt, st.session_state.chat_history_ids)
+
+        # Simulate stream of responses with milliseconds delay
+        import time
+        for chunk in response.split():
+            full_response += chunk + " "
+            time.sleep(0.05)
+            # Add a placeholder to stream the response
+            message_placeholder.markdown(full_response + "▌")
+        message_placeholder.markdown(full_response)
+
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": full_response})