Update app.py

app.py

@@ -79,15 +79,13 @@ def update_vectors(files, parser):
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
 def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2):
-    client = InferenceClient(
-        model,
-        token=huggingface_token,
-    )
-    
+    client = InferenceClient(model, token=huggingface_token)
     full_responses = []
     messages = [{"role": "user", "content": prompt}]
     
     for _ in range(num_calls):
+        if stop_clicked:  # Check if stop was clicked
+            break
         try:
             response = ""
             for message in client.chat_completion(
@@ -96,6 +94,8 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, tempe
                 temperature=temperature,
                 stream=True,
             ):
+                if stop_clicked:  # Check if stop was clicked
+                    break
                 if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                     chunk = message.choices[0].delta.content
                     response += chunk
@@ -103,8 +103,10 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, tempe
         except Exception as e:
             print(f"Error in generating response: {str(e)}")
     
-    # Combine and clean up the responses
+    # Combine all responses into a single string
     combined_response = " ".join(full_responses)
+    
+    # Clean the combined response
     clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', combined_response, flags=re.DOTALL)
     clean_response = clean_response.replace("Using the following context:", "").strip()
     clean_response = clean_response.replace("Using the following context from the PDF documents:", "").strip()
@@ -244,28 +246,29 @@ with gr.Blocks() as demo:
         clear_btn = gr.Button("Clear")
 
     def protected_generate_response(message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked):
-        if is_generating:
-            return message, history, is_generating, stop_clicked
-        is_generating = True
-        stop_clicked = False
-        
-        try:
-            if use_web_search:
-                main_content, sources = get_response_with_search(message, model, num_calls=num_calls, temperature=temperature)
-                formatted_response = f"{main_content}\n\nSources:\n{sources}"
-            else:
-                response = get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature)
-                formatted_response = response
-
-            if not stop_clicked:
-                history.append((message, formatted_response))
-        except Exception as e:
-            print(f"Error generating response: {str(e)}")
-            history.append((message, "I'm sorry, but I encountered an error while generating the response. Please try again."))
-        
-        is_generating = False
-        return "", history, is_generating, stop_clicked
-
+    if is_generating:
+        return message, history, is_generating, stop_clicked
+    is_generating = True
+    stop_clicked = False
+    
+    try:
+        if use_web_search:
+            main_content, sources = get_response_with_search(message, model, num_calls=num_calls, temperature=temperature)
+            formatted_response = f"{main_content}\n\nSources:\n{sources}"
+        else:
+            response = get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature)
+            formatted_response = response
+
+        if not stop_clicked:
+            # Only append the final, combined response to the history
+            history.append((message, formatted_response))
+    except Exception as e:
+        print(f"Error generating response: {str(e)}")
+        history.append((message, "I'm sorry, but I encountered an error while generating the response. Please try again."))
+    
+    is_generating = False
+    return "", history, is_generating, stop_clicked
+    
     submit_btn.click(
         protected_generate_response,
         inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],