Update app.py retry button

app.py

@@ -23,9 +23,6 @@ selected_model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
 client = InferenceClient(selected_model)
 
-def query_submit(user_message, history):
-    return "", history + [[user_message, None]]
-  
 def format_prompt(query, history, lookback):
     prompt = "Responses should be no more than 100 words long.\n"
     
@@ -35,6 +32,9 @@ def format_prompt(query, history, lookback):
     prompt += f"[INST] {query} [/INST]"
   
     return prompt
+    
+def query_submit(user_message, history):
+    return "", history + [[user_message, None]]
 
 def query_completion(
     query,
@@ -64,24 +64,62 @@ def query_completion(
         history[-1][1] += response.token.text
         yield history
 
+def retry_query(history):
+    history,
+    lookback = 3,
+    max_new_tokens = 256,
+):
+
+    query = history[-1][0]
+    history[-1][1]  = None
+    
+    generateKwargs = dict(
+        max_new_tokens = max_new_tokens,
+        seed = 1337,
+    )
+    
+    formatted_query = format_prompt(query, history, lookback)
+    
+    stream = client.text_generation(
+        formatted_query,
+        **generateKwargs,
+        stream = True,
+        details = True,
+        return_full_text = False
+    )
+    
+    history[-1][1] = ""
+    
+    for response in stream:
+        history[-1][1] += response.token.text
+        yield history
+
         
 """
 Chat UI using Gradio Blocks.
 
-Blocks preferred for lower-level "atomic" layout control and state management.
+Blocks preferred for "lower-level" layout control and state management.
 
 TODOs:
 - State management for dynamic components update.
-- Add scratpad readout to right of chat log.
+- Add scratchpad readout to the right of chat log.
     * Placeholder added for now.
 - Add functionality to retry button.
     * Placeholder added for now.
 - Add dropdown for model selection.
-- Add textbox for HF model selection.
+    * Placeholder added for now.
     
 """
 
 with gr.Blocks() as chatUI:
+    # gr.State()
+    
+    with gr.Row():
+        modelSelect = gr.Dropdown(
+            label = "Model selection:",
+            scale = 0.5,
+        )
+    
     with gr.Row():
         chatOutput = gr.Chatbot(
             bubble_full_width = False,
@@ -98,16 +136,11 @@ with gr.Blocks() as chatUI:
         submitButton = gr.Button("Submit", scale = 1)
         
     with gr.Row():
-        retry = gr.Button("Retry (null)")
-        clear = gr.ClearButton([queryInput, chatOutput])
-    
-    with gr.Row():
-        modelSelect = gr.Dropdown(
-            label = "Model selection:",
-        )
         fileUpload = gr.File(
             height = 100,
         )
+        retryButton = gr.Button("Retry")
+        clearButton = gr.ClearButton([queryInput, chatOutput])
     
     with gr.Row():
         with gr.Accordion(label = "Expand for edit system prompt:"):
@@ -118,11 +151,11 @@ with gr.Blocks() as chatUI:
                 scale = 4,
         )
     
-    with gr.Row():
-        footer = gr.HTML("<div class='footer'>" + selected_model + "</div>"),
-    
-        # gr.State()
+
+    """
+    Event functions
     
+    """
     queryInput.submit(
         fn = query_submit,
         inputs = [queryInput, chatOutput],
@@ -145,5 +178,12 @@ with gr.Blocks() as chatUI:
         outputs = [chatOutput],
     )
 
+    retryButton.click(
+        fn = retry_query,
+        inputs = [chatOutput],
+        outputs = [chatOutput],
+    )
+
+    
 chatUI.queue()
 chatUI.launch(show_api = False)