Update app.py

app.py

@@ -30,6 +30,9 @@ from sklearn.manifold import TSNE
 from sklearn.metrics import silhouette_score
 from scipy.stats import spearmanr
 from functools import lru_cache
+from langchain.retrievers import MultiQueryRetriever
+from langchain.llms import HuggingFacePipeline
+from transformers import pipeline
 
 # NLTK Resource Download
 def download_nltk_resources():
@@ -141,13 +144,11 @@ def preprocess_text(text, lang='german'):
 def phonetic_match(text, query, method='levenshtein_distance'):
     if method == 'levenshtein_distance':
         text_phonetic = jellyfish.soundex(text)
-        #query_phonetic = jellyfish.cologne_phonetic(query)
         query_phonetic = jellyfish.soundex(query)
         return jellyfish.levenshtein_distance(text_phonetic, query_phonetic)
     return 0
 
 def create_custom_embedding(texts, model_type='word2vec', vector_size=100, window=5, min_count=1):
-    # Tokenize the texts
     tokenized_texts = [text.split() for text in texts]
     
     if model_type == 'word2vec':
@@ -169,7 +170,6 @@ class CustomEmbeddings(HuggingFaceEmbeddings):
     def embed_query(self, text):
         return self.model.wv[text.split()]
 
-
 # Custom Tokenizer
 def create_custom_tokenizer(file_path, model_type='WordLevel', vocab_size=10000, special_tokens=None):
     with open(file_path, 'r', encoding='utf-8') as f:
@@ -191,6 +191,7 @@ def create_custom_tokenizer(file_path, model_type='WordLevel', vocab_size=10000,
     tokenizer.train_from_iterator([text], trainer)
 
     return tokenizer
+
 def custom_tokenize(text, tokenizer):
     return tokenizer.encode(text).tokens
 
@@ -220,15 +221,16 @@ def get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separator
         raise ValueError(f"Unsupported split strategy: {split_strategy}")
 
 def get_vector_store(vector_store_type, chunks, embedding_model):
-    # Convert chunks to a tuple to make it hashable
     chunks_tuple = tuple(chunks)
-    
-    # Use a helper function for the actual vector store creation
     return _create_vector_store(vector_store_type, chunks_tuple, embedding_model)
 
-
+def custom_similarity(query_embedding, doc_embedding, query, doc_text, phonetic_weight=0.3):
+    embedding_sim = np.dot(query_embedding, doc_embedding) / (np.linalg.norm(query_embedding) * np.linalg.norm(doc_embedding))
+    phonetic_sim = phonetic_match(doc_text, query)
+    combined_sim = (1 - phonetic_weight) * embedding_sim + phonetic_weight * phonetic_sim
+    return combined_sim
+	
 def _create_vector_store(vector_store_type, chunks_tuple, embedding_model):
-    # Convert the tuple back to a list for use with the vector store
     chunks = list(chunks_tuple)
     
     if vector_store_type == 'FAISS':
@@ -238,15 +240,13 @@ def _create_vector_store(vector_store_type, chunks_tuple, embedding_model):
     else:
         raise ValueError(f"Unsupported vector store type: {vector_store_type}")
 
-        
 def get_retriever(vector_store, search_type, search_kwargs):
     if search_type == 'similarity':
         return vector_store.as_retriever(search_type="similarity", search_kwargs=search_kwargs)
     elif search_type == 'mmr':
         return vector_store.as_retriever(search_type="mmr", search_kwargs=search_kwargs)
     elif search_type == 'custom':
-        # Implement custom retriever logic here
-        pass
+        return vector_store.as_retriever(search_type="similarity", search_kwargs=search_kwargs)
     else:
         raise ValueError(f"Unsupported search type: {search_type}")
 
@@ -290,15 +290,13 @@ def search_embeddings(chunks, embedding_model, vector_store_type, search_type, q
     results = sorted(results, key=score_result, reverse=True)
     end_time = time.time()
 
-    # Check if embeddings are available
     embeddings = []
     for doc in results:
         if hasattr(doc, 'embedding'):
-            embeddings.append(doc.embedding)  # Use the embedding if it exists
+            embeddings.append(doc.embedding)
         else:
-            embeddings.append(None)  # Append None if embedding doesn't exist
+            embeddings.append(None)
 
-    # Create a DataFrame with the results and embeddings
     results_df = pd.DataFrame({
         'content': [doc.page_content for doc in results],
         'embedding': embeddings
@@ -307,13 +305,12 @@ def search_embeddings(chunks, embedding_model, vector_store_type, search_type, q
     return results_df, end_time - start_time, vector_store, results
 
 # Evaluation Metrics
+# ... (previous code remains the same)
+
 def calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model, query, top_k):
     stats = {
         "num_results": len(results),
- #       "avg_content_length": sum(len(doc.page_content) for doc in results) / len(results) if results else 0,
         "avg_content_length": np.mean([len(doc.page_content) for doc in results]) if results else 0,
-
-        #"avg_content_length": np.mean([len(doc.page_content) for doc in results]) if not results.empty else 0,
         "search_time": search_time,
         "vector_store_size": vector_store._index.ntotal if hasattr(vector_store, '_index') else "N/A",
         "num_documents": len(vector_store.docstore._dict),
@@ -328,10 +325,7 @@ def calculate_statistics(results, search_time, vector_store, num_tokens, embeddi
         pairwise_similarities = np.inner(embeddings, embeddings)
         stats["result_diversity"] = 1 - np.mean(pairwise_similarities[np.triu_indices(len(embeddings), k=1)])
         
-        # Silhouette Score
         if len(embeddings) > 2:
-            print('-----')
-            #stats["silhouette_score"] = "N/A"
             stats["silhouette_score"] = silhouette_score(embeddings, range(len(embeddings)))
         else:
             stats["silhouette_score"] = "N/A"
@@ -378,24 +372,34 @@ def visualize_results(results_df, stats_df):
 def optimize_vocabulary(texts, vocab_size=10000, min_frequency=2):
     tokenizer = Tokenizer(models.BPE(unk_token="[UNK]"))
 
-    # Count word frequencies
     word_freq = Counter(word for text in texts for word in text.split())
     
-    # Remove rare words
     optimized_texts = [
         ' '.join(word for word in text.split() if word_freq[word] >= min_frequency)
         for text in texts
     ]
     
-    # Train BPE tokenizer
-#    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
     trainer = trainers.BpeTrainer(vocab_size=vocab_size, special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"])
     tokenizer.train_from_iterator(optimized_texts, trainer)
     
     return tokenizer, optimized_texts
 
+# New preprocessing function
+def optimize_query(query, llm):
+    multi_query_retriever = MultiQueryRetriever.from_llm(
+        retriever=get_retriever(vector_store, search_type, search_kwargs),
+        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])
+    return reranked_results
+
 # Main Comparison Function
-def compare_embeddings(file, query, embedding_models, custom_embedding_model, split_strategy, chunk_size, overlap_size, custom_separators, vector_store_type, search_type, top_k, lang='german', optimize_vocab=False, phonetic_weight=0.3, custom_tokenizer_file=None, custom_tokenizer_model=None, custom_tokenizer_vocab_size=10000, custom_tokenizer_special_tokens=None):
+def compare_embeddings(file, query, embedding_models, custom_embedding_model, split_strategy, chunk_size, overlap_size, custom_separators, vector_store_type, search_type, top_k, lang='german', optimize_vocab=False, phonetic_weight=0.3, custom_tokenizer_file=None, custom_tokenizer_model=None, custom_tokenizer_vocab_size=10000, custom_tokenizer_special_tokens=None, use_query_optimization=False, use_reranking=False):
     all_results = []
     all_stats = []
     settings = {
@@ -408,16 +412,16 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
         "top_k": top_k,
         "lang": lang,
         "optimize_vocab": optimize_vocab,
-        "phonetic_weight": phonetic_weight
+        "phonetic_weight": phonetic_weight,
+        "use_query_optimization": use_query_optimization,
+        "use_reranking": use_reranking
     }
 
-    # Parse embedding models
     models = [model.strip().split(':') for model in embedding_models.split(',')]
     if custom_embedding_model:
         models.append(custom_embedding_model.strip().split(':'))
 
     for model_type, model_name in models:
-        # Process the file and generate chunks & embeddings
         chunks, embedding_model, num_tokens = process_files(
             file.name if file else None,
             model_type,
@@ -433,17 +437,19 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
             custom_tokenizer_special_tokens.split(',') if custom_tokenizer_special_tokens else None
         )
 
-        # Custom embedding handling
-        #if use_custom_embedding:
-        #    custom_model = create_custom_embedding(chunks) #add custom model by name, must com from gradio FE
-        #    embedding_model = CustomEmbeddings(custom_model)
-
-        # Optimizing vocabulary if required
         if optimize_vocab:
             tokenizer, optimized_chunks = optimize_vocabulary(chunks)
             chunks = optimized_chunks
 
-        # Searching embeddings
+        if use_query_optimization:
+            llm = HuggingFacePipeline.from_model_id(
+                model_id="google/flan-t5-base",
+                task="text2text-generation",
+                model_kwargs={"temperature": 0, "max_length": 64},
+            )
+            optimized_queries = optimize_query(query, llm)
+            query = " ".join(optimized_queries)
+
         results, search_time, vector_store, results_raw = search_embeddings(
             chunks,
             embedding_model,
@@ -455,32 +461,26 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
             phonetic_weight
         )
 
-        # Storing embeddings into the results for future use
-        for doc in results_raw:
-            print(doc)  # or print(dir(doc)) to see available attributes
+        if use_reranking:
+            reranker = pipeline("text-classification", model="cross-encoder/ms-marco-MiniLM-L-12-v2")
+            results_raw = rerank_results(results_raw, query, reranker)
 
-        #embedding = doc.metadata.get('embedding', None)  # Use .get() to avoid KeyError
-
-        result_embeddings = [doc.metadata.get('embedding', None) for doc in results_raw]  # Adjust this based on the actual attribute names
-#        result_embeddings = [doc['embedding'] for doc in results_raw]  # Assuming each result has an embedding
+        result_embeddings = [doc.metadata.get('embedding', None) for doc in results_raw]
 
         stats = calculate_statistics(results_raw, search_time, vector_store, num_tokens, embedding_model, query, top_k)
         stats["model"] = f"{model_type} - {model_name}"
         stats.update(settings)
 
-        # Formatting results and attaching embeddings
         formatted_results = format_results(results_raw, stats)
         for i, result in enumerate(formatted_results):
-            result['embedding'] = result_embeddings[i]  # Add the embedding to each result
+            result['embedding'] = result_embeddings[i]
 
         all_results.extend(formatted_results)
         all_stats.append(stats)
 
-    # Create DataFrames with embeddings now included
     results_df = pd.DataFrame(all_results)
     stats_df = pd.DataFrame(all_stats)
 
-    # Visualization of the results
     fig = visualize_results(results_df, stats_df)
 
     return results_df, stats_df, fig
@@ -500,36 +500,52 @@ def format_results(results, stats):
 
 # Gradio Interface
 def launch_interface(share=True):
-    iface = gr.Interface(
-        fn=compare_embeddings,
-        inputs=[
-            gr.File(label="Upload File (Optional)"),
-            gr.Textbox(label="Search Query"),
-            gr.Textbox(label="Embedding Models (comma-separated, e.g. HuggingFace:paraphrase-miniLM,OpenAI:text-embedding-ada-002)"),
-            gr.Textbox(label="Custom Embedding Model (optional, format: type:name)"),
-            gr.Radio(choices=["token", "recursive"], label="Split Strategy", value="recursive"),
-            gr.Slider(100, 1000, step=100, value=500, label="Chunk Size"),
-            gr.Slider(0, 100, step=10, value=50, label="Overlap Size"),
-            gr.Textbox(label="Custom Split Separators (comma-separated, optional)"),
-            gr.Radio(choices=["FAISS", "Chroma"], label="Vector Store Type", value="FAISS"),
-            gr.Radio(choices=["similarity", "mmr", "custom"], label="Search Type", value="similarity"),
-            gr.Slider(1, 10, step=1, value=5, label="Top K"),
-            gr.Dropdown(choices=["german", "english", "french"], label="Language", value="german"),
-            gr.Checkbox(label="Optimize Vocabulary", value=False),
-            gr.Slider(0, 1, step=0.1, value=0.3, label="Phonetic Matching Weight"),
-            gr.File(label="Custom Tokenizer File (Optional)"),
-            gr.Textbox(label="Custom Tokenizer Model (e.g., WordLevel, BPE, Unigram)"),
-            gr.Textbox(label="Custom Tokenizer Vocab Size", value="10000"),
-            gr.Textbox(label="Custom Tokenizer Special Tokens (comma-separated)")
-        ],
-        outputs=[
-            gr.Dataframe(label="Results", interactive=False),
-            gr.Dataframe(label="Statistics", interactive=False),
-            gr.Plot(label="Visualizations")
-        ],
-        title="Advanced Embedding Comparison Tool",
-        description="Compare different embedding models and retrieval strategies with advanced preprocessing and phonetic matching"
-    )
+    with gr.Blocks() as iface:
+        gr.Markdown("# Advanced Embedding Comparison Tool")
+        
+        with gr.Tab("Simple"):
+            file_input = gr.File(label="Upload File (Optional)")
+            query_input = gr.Textbox(label="Search Query")
+            embedding_models_input = gr.Textbox(label="Embedding Models (comma-separated, e.g. HuggingFace:paraphrase-miniLM,OpenAI:text-embedding-ada-002)")
+            top_k_input = gr.Slider(1, 10, step=1, value=5, label="Top K")
+        
+        with gr.Tab("Advanced"):
+            custom_embedding_model_input = gr.Textbox(label="Custom Embedding Model (optional, format: type:name)")
+            split_strategy_input = gr.Radio(choices=["token", "recursive"], label="Split Strategy", value="recursive")
+            chunk_size_input = gr.Slider(100, 1000, step=100, value=500, label="Chunk Size")
+            overlap_size_input = gr.Slider(0, 100, step=10, value=50, label="Overlap Size")
+            custom_separators_input = gr.Textbox(label="Custom Split Separators (comma-separated, optional)")
+            vector_store_type_input = gr.Radio(choices=["FAISS", "Chroma"], label="Vector Store Type", value="FAISS")
+            search_type_input = gr.Radio(choices=["similarity", "mmr", "custom"], label="Search Type", value="similarity")
+            lang_input = gr.Dropdown(choices=["german", "english", "french"], label="Language", value="german")
+        
+        with gr.Tab("Optional"):
+            optimize_vocab_input = gr.Checkbox(label="Optimize Vocabulary", value=False)
+            phonetic_weight_input = gr.Slider(0, 1, step=0.1, value=0.3, label="Phonetic Matching Weight")
+            custom_tokenizer_file_input = gr.File(label="Custom Tokenizer File (Optional)")
+            custom_tokenizer_model_input = gr.Textbox(label="Custom Tokenizer Model (e.g., WordLevel, BPE, Unigram)")
+            custom_tokenizer_vocab_size_input = gr.Textbox(label="Custom Tokenizer Vocab Size", value="10000")
+            custom_tokenizer_special_tokens_input = gr.Textbox(label="Custom Tokenizer Special Tokens (comma-separated)")
+            use_query_optimization_input = gr.Checkbox(label="Use Query Optimization", value=False)
+            use_reranking_input = gr.Checkbox(label="Use Reranking", value=False)
+
+        results_output = gr.Dataframe(label="Results", interactive=False)
+        stats_output = gr.Dataframe(label="Statistics", interactive=False)
+        plot_output = gr.Plot(label="Visualizations")
+
+        submit_button = gr.Button("Compare Embeddings")
+        submit_button.click(
+            fn=compare_embeddings,
+            inputs=[
+                file_input, query_input, embedding_models_input, custom_embedding_model_input,
+                split_strategy_input, chunk_size_input, overlap_size_input, custom_separators_input,
+                vector_store_type_input, search_type_input, top_k_input, lang_input,
+                optimize_vocab_input, phonetic_weight_input, custom_tokenizer_file_input,
+                custom_tokenizer_model_input, custom_tokenizer_vocab_size_input,
+                custom_tokenizer_special_tokens_input, use_query_optimization_input, use_reranking_input
+            ],
+            outputs=[results_output, stats_output, plot_output]
+        )
 
     tutorial_md = """
     # Advanced Embedding Comparison Tool Tutorial
@@ -541,13 +557,10 @@ def launch_interface(share=True):
     1. Upload a file (optional) or use the default files in the system.
     2. Enter a search query.
     3. Enter embedding models as a comma-separated list (e.g., HuggingFace:paraphrase-miniLM,OpenAI:text-embedding-ada-002).
-    4. Optionally, specify a custom embedding model in the format type:name.
-    5. Choose a text splitting strategy and set chunk size and overlap.
-    6. Select a vector store type and search type.
-    7. Set the number of top results to retrieve.
-    8. Choose the language of your documents.
-    9. Optionally, optimize vocabulary or adjust phonetic matching weight.
-    10. If you have a custom tokenizer, upload the file and specify its attributes.
+    4. Set the number of top results to retrieve.
+    5. Optionally, specify advanced settings such as custom embedding models, text splitting strategies, and vector store types.
+    6. Choose whether to use optional features like vocabulary optimization, query optimization, or result reranking.
+    7. If you have a custom tokenizer, upload the file and specify its attributes.
 
     The tool will process your query and display results, statistics, and visualizations to help you compare the performance of different models and strategies.
     """
@@ -559,4 +572,5 @@ def launch_interface(share=True):
 
     iface.launch(share=share)
 
-launch_interface()
+if __name__ == "__main__":
+    launch_interface()