Update app.py

app.py

@@ -17,8 +17,19 @@ from langchain_text_splitters import (
 )
 from typing import List, Dict, Any
 import pandas as pd
-
-
+import re
+from nltk.corpus import stopwords
+from nltk.tokenize import word_tokenize
+from nltk.stem import SnowballStemmer
+import jellyfish  # For Kölner Phonetik
+from gensim.models import Word2Vec
+from gensim.models.fasttext import FastText
+from collections import Counter
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+
+nltk.download('stopwords', quiet=True)
 nltk.download('punkt', quiet=True)
 
 FILES_DIR = './files'
@@ -39,6 +50,34 @@ MODELS = {
     }
 }
 
+def preprocess_text(text, lang='german'):
+    # Convert to lowercase
+    text = text.lower()
+    
+    # Remove special characters and digits
+    text = re.sub(r'[^a-zA-Z\s]', '', text)
+    
+    # Tokenize
+    tokens = word_tokenize(text, language=lang)
+    
+    # Remove stopwords
+    stop_words = set(stopwords.words(lang))
+    tokens = [token for token in tokens if token not in stop_words]
+    
+    # Stemming
+    stemmer = SnowballStemmer(lang)
+    tokens = [stemmer.stem(token) for token in tokens]
+    
+    return ' '.join(tokens)
+
+def phonetic_match(text, query, method='koelner_phonetik'):
+    if method == 'koelner_phonetik':
+        text_phonetic = jellyfish.cologne_phonetic(text)
+        query_phonetic = jellyfish.cologne_phonetic(query)
+        return jellyfish.jaro_winkler(text_phonetic, query_phonetic)
+    # Add other phonetic methods as needed
+    return 0
+
 class FileHandler:
     @staticmethod
     def extract_text(file_path):
@@ -89,23 +128,26 @@ def get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separator
     else:
         raise ValueError(f"Unsupported split strategy: {split_strategy}")
 
-def get_vector_store(store_type, texts, embedding_model):
-    if store_type == 'FAISS':
-        return FAISS.from_texts(texts, embedding_model)
-    elif store_type == 'Chroma':
-        return Chroma.from_texts(texts, embedding_model)
+def get_vector_store(vector_store_type, chunks, embedding_model):
+    if vector_store_type == 'FAISS':
+        return FAISS.from_texts(chunks, embedding_model)
+    elif vector_store_type == 'Chroma':
+        return Chroma.from_texts(chunks, embedding_model)
     else:
-        raise ValueError(f"Unsupported vector store type: {store_type}")
+        raise ValueError(f"Unsupported vector store type: {vector_store_type}")
 
-def get_retriever(vector_store, search_type, search_kwargs=None):
+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
     else:
         raise ValueError(f"Unsupported search type: {search_type}")
 
-def process_files(file_path, model_type, model_name, split_strategy, chunk_size, overlap_size, custom_separators):
+def process_files(file_path, model_type, model_name, split_strategy, chunk_size, overlap_size, custom_separators, lang='german'):
     if file_path:
         text = FileHandler.extract_text(file_path)
     else:
@@ -113,6 +155,9 @@ def process_files(file_path, model_type, model_name, split_strategy, chunk_size,
         for file in os.listdir(FILES_DIR):
             file_path = os.path.join(FILES_DIR, file)
             text += FileHandler.extract_text(file_path)
+    
+    # Preprocess the text
+    text = preprocess_text(text, lang)
 
     text_splitter = get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separators)
     chunks = text_splitter.split_text(text)
@@ -121,15 +166,24 @@ def process_files(file_path, model_type, model_name, split_strategy, chunk_size,
 
     return chunks, embedding_model, len(text.split())
 
-def search_embeddings(chunks, embedding_model, vector_store_type, search_type, query, top_k):
+def search_embeddings(chunks, embedding_model, vector_store_type, search_type, query, top_k, lang='german', phonetic_weight=0.3):
+    # Preprocess the query
+    preprocessed_query = preprocess_text(query, lang)
+    
     vector_store = get_vector_store(vector_store_type, chunks, embedding_model)
     retriever = get_retriever(vector_store, search_type, {"k": top_k})
 
     start_time = time.time()
-    results = retriever.get_relevant_documents(query)
+    results = retriever.get_relevant_documents(preprocessed_query)
+    
+    # Apply phonetic matching
+    results = sorted(results, key=lambda x: (1 - phonetic_weight) * vector_store.similarity_search(x.page_content, k=1)[0][1] + 
+                                            phonetic_weight * phonetic_match(x.page_content, query),
+                     reverse=True)
+    
     end_time = time.time()
 
-    return results, end_time - start_time, vector_store
+    return results[:top_k], end_time - start_time, vector_store
 
 def calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model):
     return {
@@ -142,7 +196,47 @@ def calculate_statistics(results, search_time, vector_store, num_tokens, embeddi
         "embedding_vocab_size": embedding_model.client.get_vocab_size() if hasattr(embedding_model, 'client') and hasattr(embedding_model.client, 'get_vocab_size') else "N/A"
     }
 
-def compare_embeddings(file, query, model_types, model_names, split_strategy, chunk_size, overlap_size, custom_separators, vector_store_type, search_type, top_k):
+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':
+        model = Word2Vec(sentences=tokenized_texts, vector_size=vector_size, window=window, min_count=min_count, workers=4)
+    elif model_type == 'fasttext':
+        model = FastText(sentences=tokenized_texts, vector_size=vector_size, window=window, min_count=min_count, workers=4)
+    else:
+        raise ValueError("Unsupported model type")
+    
+    return model
+
+class CustomEmbeddings(HuggingFaceEmbeddings):
+    def __init__(self, model_path):
+        self.model = Word2Vec.load(model_path)  # or FastText.load() for FastText models
+    
+    def embed_documents(self, texts):
+        return [self.model.wv[text.split()] for text in texts]
+    
+    def embed_query(self, text):
+        return self.model.wv[text.split()]
+
+def optimize_vocabulary(texts, vocab_size=10000, min_frequency=2):
+    # 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 = BpeTrainer(vocab_size=vocab_size, special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"])
+    tokenizer.train_from_iterator(optimized_texts, trainer)
+    
+    return tokenizer, optimized_texts
+
+def compare_embeddings(file, query, model_types, model_names, split_strategy, chunk_size, overlap_size, custom_separators, vector_store_type, search_type, top_k, lang, use_custom_embedding, optimize_vocab, phonetic_weight):
     all_results = []
     all_stats = []
     settings = {
@@ -152,7 +246,11 @@ def compare_embeddings(file, query, model_types, model_names, split_strategy, ch
         "custom_separators": custom_separators,
         "vector_store_type": vector_store_type,
         "search_type": search_type,
-        "top_k": top_k
+        "top_k": top_k,
+        "lang": lang,
+        "use_custom_embedding": use_custom_embedding,
+        "optimize_vocab": optimize_vocab,
+        "phonetic_weight": phonetic_weight
     }
 
     for model_type, model_name in zip(model_types, model_names):
@@ -163,16 +261,27 @@ def compare_embeddings(file, query, model_types, model_names, split_strategy, ch
             split_strategy,
             chunk_size,
             overlap_size,
-            custom_separators.split(',') if custom_separators else None
+            custom_separators.split(',') if custom_separators else None,
+            lang
         )
 
+        if use_custom_embedding:
+            custom_model = create_custom_embedding(chunks)
+            embedding_model = CustomEmbeddings(custom_model)
+
+        if optimize_vocab:
+            tokenizer, optimized_chunks = optimize_vocabulary(chunks)
+            chunks = optimized_chunks
+
         results, search_time, vector_store = search_embeddings(
             chunks,
             embedding_model,
             vector_store_type,
             search_type,
             query,
-            top_k
+            top_k,
+            lang,
+            phonetic_weight
         )
 
         stats = calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model)
@@ -200,39 +309,38 @@ def format_results(results, stats):
         formatted_results.append(result)
     return formatted_results
 
-# 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.CheckboxGroup(choices=list(MODELS.keys()), label="Embedding Model Types", value=["HuggingFace"]),
-            gr.CheckboxGroup(choices=[model for models in MODELS.values() for model in models], label="Embedding Models", value=["e5-base-de"]),
+            gr.CheckboxGroup(choices=list(MODELS.keys()) + ["Custom"], label="Embedding Model Types"),
+            gr.CheckboxGroup(choices=[model for models in MODELS.values() for model in models] + ["custom_model"], label="Embedding Models"),
             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"], label="Search Type", value="similarity"),
-            gr.Slider(1, 10, step=1, value=5, label="Top K")
+            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="Use Custom Embedding", value=False),
+            gr.Checkbox(label="Optimize Vocabulary", value=False),
+            gr.Slider(0, 1, step=0.1, value=0.3, label="Phonetic Matching Weight")
         ],
         outputs=[
             gr.Dataframe(label="Results", interactive=False),
             gr.Dataframe(label="Statistics", interactive=False)
         ],
-        title="Embedding Comparison Tool",
-        description="Compare different embedding models and retrieval strategies",
-        examples=[
-            ["files/test.txt", "What is machine learning?", ["HuggingFace"], ["e5-base-de"], "recursive", 500, 50, "", "FAISS", "similarity", 5]
-        ],
-        allow_flagging="never"
+        title="Advanced Embedding Comparison Tool",
+        description="Compare different embedding models and retrieval strategies with advanced preprocessing and phonetic matching"
     )
 
     tutorial_md = """
-    # Embedding Comparison Tool Tutorial
+    # Advanced Embedding Comparison Tool Tutorial
 
-    ... (tutorial content remains the same) ...
+    ... (update the tutorial to include information about the new features) ...
     """
 
     iface = gr.TabbedInterface(