Update app.py

app.py

@@ -116,7 +116,10 @@ class FileHandler:
 def simple_tokenize(text):
     return text.split()
 
-def preprocess_text(text, lang='german'):
+def preprocess_text(text, lang='german', apply_preprocessing=True):
+    if not apply_preprocessing:
+        return text
+    
     text = text.lower()
     text = re.sub(r'[^a-zA-Z\s]', '', text)
     
@@ -141,13 +144,29 @@ def preprocess_text(text, lang='german'):
     
     return ' '.join(tokens)
 
-def phonetic_match(text, query, method='levenshtein_distance'):
+def phonetic_match(text, query, method='levenshtein_distance', apply_phonetic=True):
+    if not apply_phonetic:
+        return 0
     if method == 'levenshtein_distance':
         text_phonetic = jellyfish.soundex(text)
         query_phonetic = jellyfish.soundex(query)
         return jellyfish.levenshtein_distance(text_phonetic, query_phonetic)
     return 0
 
+def optimize_query(query, llm_model):
+    llm = HuggingFacePipeline.from_model_id(
+        model_id=llm_model,
+        task="text2text-generation",
+        model_kwargs={"temperature": 0, "max_length": 64},
+    )
+    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
+    
+
 def create_custom_embedding(texts, model_type='word2vec', vector_size=100, window=5, min_count=1):
     tokenized_texts = [text.split() for text in texts]
     
@@ -399,7 +418,7 @@ def rerank_results(results, query, reranker):
     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, use_query_optimization=False, use_reranking=False):
+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', apply_preprocessing=True, optimize_vocab=False, apply_phonetic=True, 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, query_optimization_model="google/flan-t5-base", use_reranking=False):
     all_results = []
     all_stats = []
     settings = {
@@ -431,6 +450,7 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
             overlap_size,
             custom_separators.split(',') if custom_separators else None,
             lang,
+            apply_preprocessing,
             custom_tokenizer_file,
             custom_tokenizer_model,
             int(custom_tokenizer_vocab_size),
@@ -442,12 +462,7 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
             chunks = optimized_chunks
 
         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)
+            optimized_queries = optimize_query(query, query_optimization_model)
             query = " ".join(optimized_queries)
 
         results, search_time, vector_store, results_raw = search_embeddings(
@@ -458,9 +473,10 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
             query,
             top_k,
             lang,
+            apply_phonetic,
             phonetic_weight
         )
-
+        
         if use_reranking:
             reranker = pipeline("text-classification", model="cross-encoder/ms-marco-MiniLM-L-12-v2")
             results_raw = rerank_results(results_raw, query, reranker)
@@ -506,7 +522,15 @@ def launch_interface(share=True):
         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)")
+            embedding_models_input = gr.CheckboxGroup(
+                choices=[
+                    "HuggingFace:paraphrase-miniLM",
+                    "HuggingFace:paraphrase-mpnet",
+                    "OpenAI:text-embedding-ada-002",
+                    "Cohere:embed-multilingual-v2.0"
+                ],
+                label="Embedding Models"
+            )
             top_k_input = gr.Slider(1, 10, step=1, value=5, label="Top K")
         
         with gr.Tab("Advanced"):
@@ -520,13 +544,16 @@ def launch_interface(share=True):
             lang_input = gr.Dropdown(choices=["german", "english", "french"], label="Language", value="german")
         
         with gr.Tab("Optional"):
+            apply_preprocessing_input = gr.Checkbox(label="Apply Text Preprocessing", value=True)
             optimize_vocab_input = gr.Checkbox(label="Optimize Vocabulary", value=False)
+            apply_phonetic_input = gr.Checkbox(label="Apply Phonetic Matching", value=True)
             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)
+            query_optimization_model_input = gr.Textbox(label="Query Optimization Model", value="google/flan-t5-base")
             use_reranking_input = gr.Checkbox(label="Use Reranking", value=False)
 
         results_output = gr.Dataframe(label="Results", interactive=False)
@@ -540,13 +567,15 @@ def launch_interface(share=True):
                 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
+                apply_preprocessing_input, optimize_vocab_input, apply_phonetic_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, query_optimization_model_input, use_reranking_input
             ],
             outputs=[results_output, stats_output, plot_output]
         )
 
+
     tutorial_md = """
     # Advanced Embedding Comparison Tool Tutorial