Update app.py

app.py

@@ -320,48 +320,138 @@ def phonetic_match(text, query, method='levenshtein_distance', apply_phonetic=Fa
         return jellyfish.levenshtein_distance(text_phonetic, query_phonetic)
     return 0
 
-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 = 3  # Reduce top_k for quicker test
-) -> List[str]:
-    # Initialize the language model
-    #llm = HuggingFacePipeline(pipeline(model=llm_model))
-
-    print('---- optimize query ----')
-    # Create a temporary vector store for query optimization
-    temp_vector_store = get_vector_store(vector_store_type, chunks, embedding_model)
 
-    print('---- optimize query 2 ----')
-
-    # Create a retriever with the temporary vector store
-    temp_retriever = get_retriever(temp_vector_store, search_type, {"k": top_k})
-
-    print('---- optimize query 3 ----')
-    
-    # Initialize MultiQueryRetriever with the temporary retriever and the language model
-    multi_query_retriever = MultiQueryRetriever.from_llm(
-        retriever=temp_retriever,
-        llm=llm
-    )
+from typing import List, Union
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2Gen
+import numpy as np
+from nltk.tokenize import word_tokenize
+from nltk.corpus import wordnet
+import nltk
 
-    print('---- optimize query 4 ----')
-    #print(llm.invoke('Hello'))
-    # Limit max time or set a timeout for LLM to avoid endless execution
-    try:
-        optimized_queries = multi_query_retriever.invoke(query, max_time=30)  # Timeout in seconds
-    except Exception as err:
-        print(f"Unexpected {err=}, {type(err)=}")
+def optimize_query(
+    query: str,
+    chunks: List[str],
+    embedding_model: str,
+    top_k: int = 3,
+    model_name: str = "google/flan-t5-small",  # Small model (only 80M parameters)
+    use_gpu: bool = False  # Default to CPU
+) -> str:
+    """
+    CPU-optimized version of query expansion using a small language model.
     
-    print(optimized_queries)
-    print('---- optimize query 5 ----')
+    Args:
+        query: Original search query
+        chunks: List of text chunks to search through
+        embedding_model: Name of the embedding model being used
+        top_k: Number of expansion terms to add
+        model_name: Name of the small language model to use
+        use_gpu: Whether to use GPU if available (defaults to False for CPU)
     
-    return optimized_queries
-
+    Returns:
+        Expanded query string
+    """
+    try:
+        # Set device
+        device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
+        
+        # 1. Basic text preprocessing (CPU-based)
+        tokens = word_tokenize(query.lower())
+        
+        # 2. WordNet synonyms expansion (CPU-based)
+        expanded_terms = set()
+        for token in tokens:
+            # Limit synonym lookup to save CPU resources
+            synsets = wordnet.synsets(token)[:1]  # Take only top synset per word
+            for syn in synsets:
+                # Limit number of lemmas
+                expanded_terms.update([lemma.name() for lemma in syn.lemmas()[:2]])
+        
+        # 3. Use small T5 model with reduced complexity
+        try:
+            # Load model with reduced memory footprint
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                model_max_length=128,  # Limit maximum sequence length
+                cache_dir="./model_cache"  # Cache models locally
+            )
+            model = AutoModelForSeq2Gen.from_pretrained(
+                model_name,
+                low_cpu_mem_usage=True,  # Enable low memory usage
+                device_map="cpu"  # Explicitly set to CPU
+            )
+            
+            # Move model to CPU and eval mode
+            model = model.to(device)
+            model.eval()  # Set to evaluation mode to reduce memory usage
+            
+            # Prepare input with reduced length
+            prompt = f"Enhance this search query with relevant terms: {query}"
+            inputs = tokenizer(
+                prompt,
+                return_tensors="pt",
+                max_length=64,  # Reduced from 128
+                truncation=True,
+                padding=True
+            )
+            
+            # Generate with minimal parameters
+            with torch.no_grad():  # Disable gradient calculation
+                outputs = model.generate(
+                    inputs.input_ids.to(device),
+                    max_length=32,  # Reduced from 64
+                    num_return_sequences=1,
+                    temperature=0.7,
+                    do_sample=False,  # Disable sampling for faster generation
+                    early_stopping=True
+                )
+            
+            enhanced_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Clear CUDA cache if GPU was used
+            if device == "cuda":
+                torch.cuda.empty_cache()
+                
+        except Exception as model_error:
+            print(f"Model-based expansion failed: {str(model_error)}")
+            enhanced_query = query  # Fallback to original query
+        
+        # 4. Combine original and expanded terms
+        final_terms = set(tokens)
+        final_terms.update(expanded_terms)
+        if enhanced_query != query:  # Only add if model expansion worked
+            final_terms.update(word_tokenize(enhanced_query.lower()))
+        
+        # 5. Remove stopwords and select top_k most relevant terms
+        stopwords = set(['the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to'])
+        final_terms = [term for term in final_terms if term not in stopwords]
+        
+        # Combine with original query
+        expanded_query = f"{query} {' '.join(list(final_terms)[:top_k])}"
+        
+        # Clean up
+        del model
+        del tokenizer
+        if device == "cuda":
+            torch.cuda.empty_cache()
+        print(expanded_query.strip())
+        return expanded_query.strip()
+        
+    except Exception as e:
+        print(f"Query optimization failed: {str(e)}")
+        return query  # Return original query if optimization fails
+
+# Example usage
+"""
+chunks = ["sample text chunk 1", "sample text chunk 2"]
+query = "machine learning algorithms"
+optimized_query = optimize_query(
+    query=query,
+    chunks=chunks,
+    embedding_model="sentence-transformers/all-MiniLM-L6-v2",
+    use_gpu=False  # Explicitly use CPU
+)
+"""
     
 
 def create_custom_embedding(texts, model_type='word2vec', vector_size=100, window=5, min_count=1):
@@ -781,24 +871,26 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
         if optimize_vocab:
             tokenizer, optimized_chunks = optimize_vocabulary(chunks)
             chunks = optimized_chunks
-
+        
+        search_query = query
+        
         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([doc.page_content for doc in optimized_queries])  # Extract text from Document objects
+            search_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,
             embedding_model,
             vector_store_type,
             search_type,
-            query,
+            search_query,
             top_k,
             expected_result,
             lang,
             apply_phonetic,
             phonetic_weight
-        )
+        ) 
         
         if use_reranking:
             reranker = pipeline("text-classification", model="cross-encoder/ms-marco-MiniLM-L-12-v2")