Update app.py

app.py

@@ -18,31 +18,36 @@ with open("info2.md", "r") as file:
     info2_md_content = file.read()
 
 # Chunk the info.md and info2.md content into smaller sections
-chunk_size = 2500  # Adjust this size as needed
+chunk_size = 1500  # Adjust this size as needed to fit the context window
 info_md_chunks = textwrap.wrap(info_md_content, chunk_size)
 info2_md_chunks = textwrap.wrap(info2_md_content, chunk_size)
 
 # Combine both sets of chunks
 all_chunks = info_md_chunks + info2_md_chunks
 
+# Function to initialize the conversation history with chunks
+def initialize_history(chunks):
+    history = []
+    for chunk in chunks:
+        history.append(("System Information", chunk))
+    return history
+
+# Initialize history with initial chunks
+history = initialize_history(all_chunks[:2])  # Starting with the first two chunks for example
+
 def format_prompt_mixtral(message, history, chunks):
     prompt = "<s>"
     prompt += f"{system_prompt_text}\n\n"  # Add the system prompt
 
     # Include the initial context from the chunks
-    for chunk in chunks:
-        prompt += f"[INST] System Information [/INST] {chunk}</s> "
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST] {bot_response}</s> "
 
-    # Add conversation history
-    if history:
-        for user_prompt, bot_response in history:
-            prompt += f"[INST] {user_prompt} [/INST] {bot_response}</s> "
-    
     # Add the current user message
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def chat_inf(message, history, seed, temp, tokens, top_p, rep_p):
+def chat_inf(message, history, chunks, seed, temp, tokens, top_p, rep_p):
     generate_kwargs = dict(
         temperature=temp,
         max_new_tokens=tokens,
@@ -52,7 +57,7 @@ def chat_inf(message, history, seed, temp, tokens, top_p, rep_p):
         seed=seed,
     )
 
-    formatted_prompt = format_prompt_mixtral(message, history, all_chunks)
+    formatted_prompt = format_prompt_mixtral(message, history, chunks)
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
     for response in stream:
@@ -62,6 +67,8 @@ def chat_inf(message, history, seed, temp, tokens, top_p, rep_p):
     yield history
 
 def clear_fn():
+    global history
+    history = initialize_history(all_chunks[:2])  # Reset to initial chunks
     return None, None
 
 rand_val = random.randint(1, 1111111111111111)
@@ -72,7 +79,7 @@ def check_rand(inp, val):
     else:
         return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=int(val))
 
-with gr.Blocks() as app:  # Add auth here
+with gr.Blocks() as app:
     gr.HTML("""<center><h1 style='font-size:xx-large;'>PTT Chatbot</h1><br><h3>running on Huggingface Inference </h3><br><h7>EXPERIMENTAL</center>""")
     with gr.Row():
         chat = gr.Chatbot(height=500)
@@ -98,7 +105,7 @@ with gr.Blocks() as app:  # Add auth here
 
     hid1 = gr.Number(value=1, visible=False)
 
-    go = btn.click(check_rand, [rand, seed], seed).then(chat_inf, [inp, chat, seed, temp, tokens, top_p, rep_p], chat)
+    go = btn.click(check_rand, [rand, seed], seed).then(chat_inf, [inp, chat, all_chunks, seed, temp, tokens, top_p, rep_p], chat)
 
     stop_btn.click(None, None, None, cancels=[go])
     clear_btn.click(clear_fn, None, [inp, chat])