Update app.py

app.py

@@ -79,7 +79,8 @@ def generate_chunked_response(prompt, max_tokens=1000, max_chunks=5):
             "temperature": 0.7,
             "top_p": 0.95,
             "top_k": 40,
-            "repetition_penalty": 1.1
+            "repetition_penalty": 1.1,
+            "stop": ["</s>", "[/INST]"]  # Add stop tokens
         }
     }
     
@@ -90,14 +91,31 @@ def generate_chunked_response(prompt, max_tokens=1000, max_chunks=5):
             result = response.json()
             if isinstance(result, list) and len(result) > 0:
                 chunk = result[0].get('generated_text', '')
-                full_response += chunk
-                if chunk.endswith((".", "!", "?")):
+                
+                # Remove any part of the chunk that's already in full_response
+                new_content = chunk[len(full_response):].strip()
+                
+                if not new_content:
+                    break  # No new content, so we're done
+                
+                full_response += new_content
+                
+                if chunk.endswith((".", "!", "?", "</s>", "[/INST]")):
                     break
+                
+                # Update the prompt for the next iteration
+                payload["inputs"] = full_response
             else:
                 break
         else:
             break
-    return full_response.strip()
+    
+    # Clean up the response
+    clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', full_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()
+    
+    return clean_response
 
 def duckduckgo_search(query):
     with DDGS() as ddgs: