add chat

app.py

@@ -17,6 +17,23 @@ def load():
         )
     return model, tokenizer
 
+def get_prompt(user_query, system_prompt, messages="", sep="\n\n### "):
+    prompt = system_prompt + "\n以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。"
+    roles = ["指示", "応答"]
+    msgs = [": \n" + user_query, ": "]
+    if messages:
+        roles.insert(1, "入力")
+        msgs.insert(1, ": \n" + "\n".join(message for _, message in messages))
+
+    for role, msg in zip(roles, msgs):
+        prompt += sep + role + msg
+    return prompt
+
+def get_input_token_length(user_query, system_prompt, messages=""):
+    prompt = get_prompt(user_query, system_prompt, messages)
+    input_ids = tokenizer([prompt], return_tensors='np', add_special_tokens=False)['input_ids']
+    return input_ids.shape[-1]
+
 def generate():
     pass
 
@@ -25,6 +42,7 @@ st.header(":dna: 遺伝カウンセリング対話AI")
 
 
 # 初期化
+model, tokenizer = load()
 if "messages" not in st.session_state:
     st.session_state["messages"] = []
 if "options" not in st.session_state:
@@ -64,9 +82,9 @@ if user_prompt := st.chat_input("質問を送信してください"):
     with st.chat_message("user"):
         st.text(user_prompt)
     st.session_state["messages"].append({"role": "user", "content": user_prompt})
-    response = None
+    token_kength = get_input_token_length(user_query=user_prompt, system_prompt=st.session_state["options"]["system_prompt"], messages=st.session_state["messages"])
+    response = f"{token_kength}: " + get_prompt(user_query=user_prompt, system_prompt=st.session_state["options"]["system_prompt"], messages=st.session_state["messages"])
     with st.chat_message("assistant"):
         st.text(response)
     st.session_state["messages"].append({"role": "assistant", "content": user_prompt})
 
-model, tokenizer = load()