Reverted back to the initial app

app.py

@@ -1,6 +1,6 @@
+from openai import OpenAI
 import streamlit as st
 import os
-from openai import OpenAI
 
 api_key = os.getenv("NVIDIANIM_API_KEY")
 
@@ -11,41 +11,29 @@ client = OpenAI(
 
 model_name = "meta/llama-3.1-405b-instruct"
 
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
 def get_llama_response(question):
-    st.session_state.messages.append({"role": "user", "content": question})
-    response = client.chat.completions.create(
-        model=model_name,
-        messages=st.session_state.messages,
-        temperature=0.2,
-        top_p=0.7,
-        max_tokens=1024,
-        stream=True
-    )
-
-    response_text = ""
-    for chunk in response:
-        if chunk.choices[0].delta.content is not None:
-            response_text += chunk.choices[0].delta.content
-
-    st.session_state.messages.append({"role": "assistant", "content": response_text})
-    return response_text
+  completion = client.chat.completions.create(
+      model=model_name,
+      messages=[{"role": "user", "content": question}],
+      temperature=0.2,
+      top_p=0.7,
+      max_tokens=1024,
+      stream=True
+  )
+  
+  response = ""
+  for chunk in completion:
+    if chunk.choices[0].delta.content is not None:
+      response += chunk.choices[0].delta.content
+  return response.strip()
 
 st.title("Ask Llama 3.1 405B on Nvidia NIM")
-
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-user_input = st.chat_input("Your message")
+user_question = st.text_input("Enter your question:")
 
 if st.button("Submit"):
-    if user_input:
-        st.session_state.messages.append({"role": "user", "content": user_input})
-        with st.chat_message("user"):
-            st.markdown(user_input)
-        response = get_llama_response(user_input)
-        with st.chat_message("assistant"):
-            st.markdown(response)
+  if user_question:
+    llama_response = get_llama_response(user_question)
+    st.write("**Llama 3.1 405B Response:**")
+    st.write(llama_response)
+  else:
+    st.warning("Please enter a question.")