Update app.py

app.py

@@ -46,9 +46,9 @@ login(token=hf_token)
 # Define the model pipeline with additional generation parameters
 model_pipeline = pipeline(
    # model="meta-llama/Llama-3.2-1B",
-    model="meta-llama/Llama-3.2-1B",
+    model="meta-llama/Llama-3.2-3B-Instruct",
     #use_auth_token=hf_token,
-    max_length=1000,  # You can increase this if needed
+    #max_length=1000,  # You can increase this if needed
     max_new_tokens=500  # Limit how many tokens are generated
 )
 
@@ -177,31 +177,33 @@ def phonetic_match(text, query, method='levenshtein_distance', apply_phonetic=Tr
     return 0
 
 #def optimize_query(query, llm_model):
-def optimize_query(query, llm_model, chunks, embedding_model, vector_store_type, search_type, top_k):
-    # Use a HuggingFace model for text generation
-    #model_id = "google/flan-t5-large"
-    #tokenizer = AutoTokenizer.from_pretrained(model_id)
-    #model = AutoModelForCausalLM.from_pretrained(model_id)
-    #pipe = pipeline(
-    #    "text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512
-    #)
-    #llm = HuggingFacePipeline(pipeline=pipe)
+def optimize_query(
+    query: str, 
+    llm_model: str = "meta-llama/Llama-3.2-1B", 
+    chunks: List[str] = None, 
+    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2", 
+    vector_store_type: str = "faiss", 
+    search_type: str = "similarity", 
+    top_k: int = 5
+) -> List[str]:
+    # Initialize the language model
+    llm = HuggingFacePipeline(model=llm_model)
 
-    #llm = HuggingFacePipeline(pipeline(model="HuggingFaceH4/zephyr-7b-beta"))
-    
-    
     # 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})
-    
+
+    # Initialize MultiQueryRetriever with the temporary retriever and the language model
     multi_query_retriever = MultiQueryRetriever.from_llm(
         retriever=temp_retriever,
         llm=llm
-    ) 
-        # Use a NoOpRunManager as the run_manager 
+    )
+
+    # Use a NoOpRunManager as the run manager
     optimized_queries = multi_query_retriever.invoke(query)
+
     return optimized_queries 
     
 
@@ -593,7 +595,8 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
 
         if use_query_optimization:
             optimized_queries = optimize_query(query, query_optimization_model, chunks, embedding_model, vector_store_type, search_type, top_k)
-            query = " ".join(optimized_queries)
+            #query = " ".join(optimized_queries)
+            query = " ".join([doc.page_content for doc in optimized_queries])  # Extract text from Document objects
 
         results, search_time, vector_store, results_raw = search_embeddings(
             chunks,