Update app.py

app.py

@@ -51,10 +51,10 @@ def format_prompt(message, history):
     prompt += f"<start_of_turn>user{message}<end_of_turn><start_of_turn>model"
     #print(prompt)
     return prompt
+result = []
 
 
-
-def chat_inf(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,rep_p):
+def chat_inf(system_prompt,prompt,history,memory,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem):
     #token max=8192
     hist_len=0
     client=clients[int(client_choice)-1]
@@ -62,17 +62,19 @@ def chat_inf(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,r
         history = []
         hist_len=0
     if history:
-        for ea in history:
+        for ea in history[0-chat_mem:]:
             hist_len+=len(str(ea))
         print(hist_len)
     in_len=len(system_prompt+prompt)+hist_len
+    
+    
     print("\n######### HIST "+str(in_len))
     print("\n######### TOKENS "+str(tokens))
     if (in_len+tokens) > 8000:
         yield [(prompt,"Wait. I need to compress our Chat history...")]
-        history=compress_history(history,client_choice,seed,temp,tokens,top_p,rep_p)
+        hist=compress_history(history[-5:],client_choice,seed,temp,tokens,top_p,rep_p)
         yield [(prompt,"History has been compressed, processing request...")]
-            
+        history = [(prompt,hist)]
     generate_kwargs = dict(
         temperature=temp,
         max_new_tokens=tokens,
@@ -82,7 +84,7 @@ def chat_inf(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,r
         seed=seed,
     )
     #formatted_prompt=prompt   
-    formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+    formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history[0-chat_mem:])
     print("\n######### PROMPT "+str(len(formatted_prompt)))
 
 
@@ -95,7 +97,8 @@ def chat_inf(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,r
         output += response.token.text
         yield [(prompt,output)]
     history.append((prompt,output))
-    yield history
+    memory=history
+    yield history,memory
 
 def get_screenshot(chat: list,height=5000,width=600,chatblock=[],theme="light",wait=3000,header=True):
     print(chatblock)
@@ -121,6 +124,7 @@ def check_rand(inp,val):
 
     
 with gr.Blocks() as app:
+    memory=gr.State()
     gr.HTML("""<center><h1 style='font-size:xx-large;'>Google Gemma Models</h1><br><h3>running on Huggingface Inference Client</h3><br><h7>EXPERIMENTAL""")
     chat_b = gr.Chatbot(height=500)
     with gr.Group():
@@ -145,6 +149,7 @@ with gr.Blocks() as app:
                     temp=gr.Slider(label="Temperature",step=0.01, minimum=0.01, maximum=1.0, value=0.9)
                     top_p=gr.Slider(label="Top-P",step=0.01, minimum=0.01, maximum=1.0, value=0.9)
                     rep_p=gr.Slider(label="Repetition Penalty",step=0.1, minimum=0.1, maximum=2.0, value=1.0)
+                    chat_mem=gr.Number(label="Chat Memory", info="Number of previous chats to retain",value=5)
         with gr.Accordion(label="Screenshot",open=False):
             with gr.Row():
                 with gr.Column(scale=3):
@@ -161,8 +166,10 @@ with gr.Blocks() as app:
 
 
     im_go=im_btn.click(get_screenshot,[chat_b,im_height,im_width,chatblock,theme,wait_time],img)
-    chat_sub=inp.submit(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p],chat_b)
-    go=btn.click(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p],chat_b)
+    chat_sub=inp.submit(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem],chat_b)
+    
+    go=btn.click(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,memory,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem],[chat_b,memory])
+    
     stop_btn.click(None,None,None,cancels=[go,im_go,chat_sub])
     clear_btn.click(clear_fn,None,[inp,sys_inp,chat_b])
 app.queue(default_concurrency_limit=10).launch()