Update app.py

app.py

@@ -14,7 +14,14 @@ retriever = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 def generate_query(document):
     prompt = f"Generate a relevant search query for the following document:\n\n{document}\n\nQuery:"
     input_ids = llm_tokenizer.encode(prompt, return_tensors="pt")
-    output = llm.generate(input_ids, max_length=50, num_return_sequences=5)
+    output = llm.generate(
+        input_ids, 
+        max_length=50, 
+        num_return_sequences=5,
+        num_beams=5,  # Use beam search
+        no_repeat_ngram_size=2,
+        early_stopping=True
+    )
     queries = [llm_tokenizer.decode(seq, skip_special_tokens=True) for seq in output]
     return queries