Update app.py

app.py

@@ -434,14 +434,26 @@ def optimize_vocabulary(texts, vocab_size=10000, min_frequency=2):
     return tokenizer, optimized_texts
 
 # New preprocessing function
-def optimize_query(query, llm):
+def optimize_query(query, llm_model, chunks, embedding_model, vector_store_type, search_type, top_k):
+    llm = HuggingFacePipeline.from_model_id(
+        model_id=llm_model,
+        task="text2text-generation",
+        model_kwargs={"do_sample": True, "temperature": 0, "max_new_tokens": 64},
+    )
+    
+    # Create a temporary vector store for query optimization
+    temp_vector_store = get_vector_store(vector_store_type, chunks, embedding_model)
+    
+    # Create a retriever with the temporary vector store
+    temp_retriever = get_retriever(temp_vector_store, search_type, {"k": top_k})
+    
     multi_query_retriever = MultiQueryRetriever.from_llm(
-        retriever=get_retriever(vector_store, search_type, search_kwargs),
+        retriever=temp_retriever,
         llm=llm
     )
     optimized_queries = multi_query_retriever.generate_queries(query)
     return optimized_queries
-
+    
 # New postprocessing function
 def rerank_results(results, query, reranker):
     reranked_results = reranker.rerank(query, [doc.page_content for doc in results])
@@ -495,7 +507,7 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
             chunks = optimized_chunks
 
         if use_query_optimization:
-            optimized_queries = optimize_query(query, query_optimization_model)
+            optimized_queries = optimize_query(query, query_optimization_model, chunks, embedding_model, vector_store_type, search_type, top_k)
             query = " ".join(optimized_queries)
 
         results, search_time, vector_store, results_raw = search_embeddings(