Update app.py

app.py

@@ -46,12 +46,23 @@ def legal_assistant_query(query):
     query_embedding = embedder.encode([query])
     D, I = index.search(np.array(query_embedding), k=5)
 
+    # Limit the number of retrieved documents or trim context
     retrieved_docs = [corpus[i] for i in I[0]]
-    context_combined = "\n\n".join(retrieved_docs)
+    
+    # Combine retrieved documents into a single context and ensure it doesn't exceed token limit
+    context_combined = "\n\n".join(retrieved_docs[:3])  # Limit to 3 docs to avoid overflow
+    max_length = 1024  # Set appropriate limit based on GPT-2's token length (around 1024 tokens)
+    
+    # Ensure the context combined does not exceed max length
+    context_combined = context_combined[:max_length]
 
+    # Prepare the prompt for GPT-2
     prompt = f"Given the following legal references, answer the question:\n\n{context_combined}\n\nQuestion: {query}\nAnswer:"
+
+    # Generate the response
     result = generator(prompt, max_new_tokens=200, do_sample=True)[0]['generated_text']
     
+    # Extract the answer from the generated text
     return result.split("Answer:")[-1].strip()
 
 # Gradio Interface