Update app.py

app.py

@@ -1,35 +1,44 @@
 import json 
 import os
 
+import openai
 import requests
 import streamlit as st
 
 CHATBOT_ENDPOINT = os.environ["CHATBOT_ENDPOINT"]
 TOKEN = os.environ["TOKEN"]
+MAINTENANCE = os.environ.get("MAINTENANCE", 0)
+
+client = openai.OpenAI(
+    base_url=CHATBOT_ENDPOINT,
+    api_key=TOKEN,
+    timeout=60
+)
 
 def generate(prompt):
 
     try:
-        r = requests.post(
-            CHATBOT_ENDPOINT,
-            data=json.dumps({"instruction": prompt, "token": TOKEN}),
-            headers = {"content-type": "application/json"},
-            timeout = 20,
-            stream = True,
+        completion = client.chat.completions.create(
+            model="stockmark/stockmark-100b-instruct-merged-v0.1",
+            messages=[{"role": "user", "content": prompt}],
+            extra_body={"repetition_penalty": 1.05},
+            max_tokens=256,
+            temperature=0.5,
+            top_p=0.95,
+            stream=True
         )
-        new_str = b""
-        for s in r.iter_content():
-            new_str += s
-            try:
-                output_str = new_str.decode("utf-8")
-                yield output_str.replace("\n", "  \n")
-                new_str = b""
-            except:
-                pass
+        
+        for s in completion:
+            s = s.choices[0].delta.content
+            if s:
+                if s == "\n":
+                    s = "  \n"
+                yield s
     except:
-        yield "<<Some errors occured>>"
+        yield "<<予期せぬエラーが発生しております。しばらくしてからアクセスください。>>"
+
 
-intro = """This is a demo site for Stockmark-LLM-100b. This service is running on AWS Inferentia2.
+intro = """This is a demo site for Stockmark-LLM-100b. This service is running on AWS Inferentia2. Currently, the response is sometimes slow due to many requests to the server.
 - Pretrained model: [stockmark/stockmark-100b](https://huggingface.co/stockmark/stockmark-100b)
 - Instruction tuned model: [stockmark/stockmark-100b-instruct-v0.1](https://huggingface.co/stockmark/stockmark-100b-instruct-v0.1)
 """
@@ -39,6 +48,42 @@ disclaimer = """
 - We may use users chat data in this demo to improve our LLM.
 """
 
-st.title("Stockmark-LLM-100b")    
-st.markdown("申し訳ありませんが、ただいまサーバの復旧をおこなっております。しばらくしてからアクセスしてください。")
+if MAINTENANCE:
+    st.title("Stockmark-LLM-100b") 
+    st.markdown("ただいまメンテナンス中です。申し訳ありませんが、しばらくしてからアクセスしてください。")
+    st.stop()
+
+tab1, tab2 = st.tabs(["Demo", "Disclaimer"])
+
+with tab1:
+    st.title("Stockmark-LLM-100b")    
+    st.markdown(intro)
+    
+    prompt = st.session_state.get("prompt", "")
+    response = st.session_state.get("response", "")
+    
+    if prompt == "" or response:
+        print("new_session")
+        prompt_new = st.text_area("Prompt:")
+        if prompt_new:
+            st.session_state["prompt"] = prompt_new
+            st.session_state["response"] = ""
+            st.rerun()
+    else:
+        prompt = st.text_area("Prompt:", value=prompt, disabled=True)
+    
+    if prompt:
+    
+        if response:
+            with st.chat_message("assistant"):
+                st.write(response)
+        else:
+            with st.chat_message("assistant"):
+                response = st.write_stream(generate(prompt))
+            
+            st.session_state["response"] = response
+            st.rerun()
 
+with tab2:
+    st.title("Stockmark-LLM-100b: Disclaimer")
+    st.markdown(disclaimer)