Update app.py

app.py

@@ -5,38 +5,35 @@ import docx
 import nltk
 import gradio as gr
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.embeddings import (
-    OpenAIEmbeddings,
-    CohereEmbeddings,
-)
+from langchain_community.embeddings import CohereEmbeddings
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import FAISS, Chroma
-from langchain_text_splitters import (
-    RecursiveCharacterTextSplitter,
-    TokenTextSplitter,
-)
+from langchain_text_splitters import RecursiveCharacterTextSplitter, TokenTextSplitter
 from typing import List, Dict, Any
 import pandas as pd
+import numpy as np
 import re
 from nltk.corpus import stopwords
 from nltk.tokenize import word_tokenize
 from nltk.stem import SnowballStemmer
-import jellyfish  # For Kölner Phonetik
+import jellyfish
 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
-
-
+from tokenizers.models import WordLevel
+from tokenizers.trainers import WordLevelTrainer
+from tokenizers.pre_tokenizers import Whitespace
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.manifold import TSNE
+from sklearn.metrics import silhouette_score
+from scipy.stats import spearmanr
+from functools import lru_cache
 
+# NLTK Resource Download
 def download_nltk_resources():
-    resources = [
-        'punkt',
-        'stopwords',
-        'snowball_data',
-    ]
+    resources = ['punkt', 'stopwords', 'snowball_data']
     for resource in resources:
         try:
             nltk.download(resource, quiet=True)
@@ -45,47 +42,87 @@ def download_nltk_resources():
 
 download_nltk_resources()
 
+FILES_DIR = './files'
 
-nltk.download('stopwords', quiet=True)
-nltk.download('punkt', quiet=True)
+# Model Management
+class ModelManager:
+    def __init__(self):
+        self.models = {
+            'HuggingFace': {
+                'e5-base-de': "danielheinz/e5-base-sts-en-de",
+                'paraphrase-miniLM': "paraphrase-multilingual-MiniLM-L12-v2",
+                'paraphrase-mpnet': "paraphrase-multilingual-mpnet-base-v2",
+                'gte-large': "gte-large",
+                'gbert-base': "gbert-base"
+            },
+            'OpenAI': {
+                'text-embedding-ada-002': "text-embedding-ada-002"
+            },
+            'Cohere': {
+                'embed-multilingual-v2.0': "embed-multilingual-v2.0"
+            }
+        }
 
-FILES_DIR = './files'
+    def add_model(self, provider, name, model_path):
+        if provider not in self.models:
+            self.models[provider] = {}
+        self.models[provider][name] = model_path
 
-MODELS = {
-    'HuggingFace': {
-        'e5-base-de': "danielheinz/e5-base-sts-en-de",
-        'paraphrase-miniLM': "paraphrase-multilingual-MiniLM-L12-v2",
-        'paraphrase-mpnet': "paraphrase-multilingual-mpnet-base-v2",
-        'gte-large': "gte-large",
-        'gbert-base': "gbert-base"
-    },
-    'OpenAI': {
-        'text-embedding-ada-002': "text-embedding-ada-002"
-    },
-    'Cohere': {
-        'embed-multilingual-v2.0': "embed-multilingual-v2.0"
-    }
-}
+    def remove_model(self, provider, name):
+        if provider in self.models and name in self.models[provider]:
+            del self.models[provider][name]
+
+    def get_model(self, provider, name):
+        return self.models.get(provider, {}).get(name)
+
+    def list_models(self):
+        return {provider: list(models.keys()) for provider, models in self.models.items()}
+
+model_manager = ModelManager()
 
+# File Handling
+class FileHandler:
+    @staticmethod
+    def extract_text(file_path):
+        ext = os.path.splitext(file_path)[-1].lower()
+        if ext == '.pdf':
+            return FileHandler._extract_from_pdf(file_path)
+        elif ext == '.docx':
+            return FileHandler._extract_from_docx(file_path)
+        elif ext == '.txt':
+            return FileHandler._extract_from_txt(file_path)
+        else:
+            raise ValueError(f"Unsupported file type: {ext}")
+
+    @staticmethod
+    def _extract_from_pdf(file_path):
+        with pdfplumber.open(file_path) as pdf:
+            return ' '.join([page.extract_text() for page in pdf.pages])
+
+    @staticmethod
+    def _extract_from_docx(file_path):
+        doc = docx.Document(file_path)
+        return ' '.join([para.text for para in doc.paragraphs])
+
+    @staticmethod
+    def _extract_from_txt(file_path):
+        with open(file_path, 'r', encoding='utf-8') as f:
+            return f.read()
+
+# Text Processing
 def simple_tokenize(text):
-    """Simple tokenization fallback method."""
     return text.split()
 
 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
     try:
         tokens = word_tokenize(text, language=lang)
     except LookupError:
         print(f"Warning: NLTK punkt tokenizer for {lang} not found. Using simple tokenization.")
         tokens = simple_tokenize(text)
     
-    # Remove stopwords
     try:
         stop_words = set(stopwords.words(lang))
     except LookupError:
@@ -93,7 +130,6 @@ def preprocess_text(text, lang='german'):
         stop_words = set()
     tokens = [token for token in tokens if token not in stop_words]
     
-    # Stemming
     try:
         stemmer = SnowballStemmer(lang)
         tokens = [stemmer.stem(token) for token in tokens]
@@ -107,44 +143,34 @@ def phonetic_match(text, query, 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):
-        ext = os.path.splitext(file_path)[-1].lower()
-        if ext == '.pdf':
-            return FileHandler._extract_from_pdf(file_path)
-        elif ext == '.docx':
-            return FileHandler._extract_from_docx(file_path)
-        elif ext == '.txt':
-            return FileHandler._extract_from_txt(file_path)
-        else:
-            raise ValueError(f"Unsupported file type: {ext}")
+# Custom Tokenizer
+def create_custom_tokenizer(file_path):
+    with open(file_path, 'r', encoding='utf-8') as f:
+        text = f.read()
 
-    @staticmethod
-    def _extract_from_pdf(file_path):
-        with pdfplumber.open(file_path) as pdf:
-            return ' '.join([page.extract_text() for page in pdf.pages])
+    tokenizer = Tokenizer(WordLevel(unk_token="[UNK]"))
+    tokenizer.pre_tokenizer = Whitespace()
 
-    @staticmethod
-    def _extract_from_docx(file_path):
-        doc = docx.Document(file_path)
-        return ' '.join([para.text for para in doc.paragraphs])
+    trainer = WordLevelTrainer(special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"])
+    tokenizer.train_from_iterator([text], trainer)
 
-    @staticmethod
-    def _extract_from_txt(file_path):
-        with open(file_path, 'r', encoding='utf-8') as f:
-            return f.read()
+    return tokenizer
 
+def custom_tokenize(text, tokenizer):
+    return tokenizer.encode(text).tokens
+
+# Embedding and Vector Store
+@lru_cache(maxsize=None)
 def get_embedding_model(model_type, model_name):
+    model_path = model_manager.get_model(model_type, model_name)
     if model_type == 'HuggingFace':
-        return HuggingFaceEmbeddings(model_name=MODELS[model_type][model_name])
+        return HuggingFaceEmbeddings(model_name=model_path)
     elif model_type == 'OpenAI':
-        return OpenAIEmbeddings(model=MODELS[model_type][model_name])
+        return OpenAIEmbeddings(model=model_path)
     elif model_type == 'Cohere':
-        return CohereEmbeddings(model=MODELS[model_type][model_name])
+        return CohereEmbeddings(model=model_path)
     else:
         raise ValueError(f"Unsupported model type: {model_type}")
 
@@ -160,6 +186,7 @@ def get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separator
     else:
         raise ValueError(f"Unsupported split strategy: {split_strategy}")
 
+@lru_cache(maxsize=None)
 def get_vector_store(vector_store_type, chunks, embedding_model):
     if vector_store_type == 'FAISS':
         return FAISS.from_texts(chunks, embedding_model)
@@ -179,7 +206,8 @@ def get_retriever(vector_store, search_type, search_kwargs):
     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, lang='german'):
+# Main Processing Functions
+def process_files(file_path, model_type, model_name, split_strategy, chunk_size, overlap_size, custom_separators, lang='german', custom_tokenizer_file=None):
     if file_path:
         text = FileHandler.extract_text(file_path)
     else:
@@ -188,8 +216,11 @@ def process_files(file_path, model_type, model_name, split_strategy, chunk_size,
             file_path = os.path.join(FILES_DIR, file)
             text += FileHandler.extract_text(file_path)
     
-    # Preprocess the text
-    text = preprocess_text(text, lang)
+    if custom_tokenizer_file:
+        tokenizer = create_custom_tokenizer(custom_tokenizer_file)
+        text = ' '.join(custom_tokenize(text, tokenizer))
+    else:
+        text = preprocess_text(text, lang)
 
     text_splitter = get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separators)
     chunks = text_splitter.split_text(text)
@@ -199,7 +230,6 @@ 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, 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)
@@ -208,15 +238,18 @@ def search_embeddings(chunks, embedding_model, vector_store_type, search_type, q
     start_time = time.time()
     results = retriever.invoke(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)
+    def score_result(doc):
+        similarity_score = vector_store.similarity_search_with_score(doc.page_content, k=1)[0][1]
+        phonetic_score = phonetic_match(doc.page_content, query)
+        return (1 - phonetic_weight) * similarity_score + phonetic_weight * phonetic_score
+
+    results = sorted(results, key=score_result, reverse=True)
     
     end_time = time.time()
 
     return results[:top_k], end_time - start_time, vector_store
 
+# Evaluation Metrics
 def calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model, query, top_k):
     stats = {
         "num_results": len(results),
@@ -230,64 +263,58 @@ def calculate_statistics(results, search_time, vector_store, num_tokens, embeddi
         "top_k": top_k,
     }
     
-    # Calculate diversity of results
     if len(results) > 1:
         embeddings = [embedding_model.embed_query(doc.page_content) for doc in results]
-        pairwise_similarities = cosine_similarity(embeddings)
+        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:
+            stats["silhouette_score"] = silhouette_score(embeddings, range(len(embeddings)))
+        else:
+            stats["silhouette_score"] = "N/A"
     else:
         stats["result_diversity"] = "N/A"
+        stats["silhouette_score"] = "N/A"
     
-    # Calculate rank correlation between embedding similarity and result order
     query_embedding = embedding_model.embed_query(query)
     result_embeddings = [embedding_model.embed_query(doc.page_content) for doc in results]
-    similarities = [cosine_similarity([query_embedding], [emb])[0][0] for emb in result_embeddings]
+    similarities = [np.inner(query_embedding, emb)[0] for emb in result_embeddings]
     rank_correlation, _ = spearmanr(similarities, range(len(similarities)))
     stats["rank_correlation"] = rank_correlation
     
     return stats
 
-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
+# Visualization
+def visualize_results(results_df, stats_df):
+    fig, axs = plt.subplots(2, 2, figsize=(20, 20))
     
-    def embed_documents(self, texts):
-        return [self.model.wv[text.split()] for text in texts]
+    sns.barplot(x='model', y='search_time', data=stats_df, ax=axs[0, 0])
+    axs[0, 0].set_title('Search Time by Model')
+    axs[0, 0].set_xticklabels(axs[0, 0].get_xticklabels(), rotation=45, ha='right')
     
-    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())
+    sns.scatterplot(x='result_diversity', y='rank_correlation', hue='model', data=stats_df, ax=axs[0, 1])
+    axs[0, 1].set_title('Result Diversity vs. Rank Correlation')
     
-    # Remove rare words
-    optimized_texts = [
-        ' '.join(word for word in text.split() if word_freq[word] >= min_frequency)
-        for text in texts
-    ]
+    sns.boxplot(x='model', y='avg_content_length', data=stats_df, ax=axs[1, 0])
+    axs[1, 0].set_title('Distribution of Result Content Lengths')
+    axs[1, 0].set_xticklabels(axs[1, 0].get_xticklabels(), rotation=45, ha='right')
     
-    # 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)
+    embeddings = np.array([embedding for embedding in results_df['embedding'] if isinstance(embedding, np.ndarray)])
+    if len(embeddings) > 1:
+        tsne = TSNE(n_components=2, random_state=42)
+        embeddings_2d = tsne.fit_transform(embeddings)
+        
+        sns.scatterplot(x=embeddings_2d[:, 0], y=embeddings_2d[:, 1], hue=results_df['model'][:len(embeddings)], ax=axs[1, 1])
+        axs[1, 1].set_title('t-SNE Visualization of Result Embeddings')
+    else:
+        axs[1, 1].text(0.5, 0.5, "Not enough data for t-SNE visualization", ha='center', va='center')
     
-    return tokenizer, optimized_texts
+    plt.tight_layout()
+    return fig
 
-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='german', use_custom_embedding=False, optimize_vocab=False, phonetic_weight=0.3):
+# Main Comparison Function
+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='german', use_custom_embedding=False, optimize_vocab=False, phonetic_weight=0.3, custom_tokenizer_file=None):
     all_results = []
     all_stats = []
     settings = {
@@ -313,7 +340,8 @@ def compare_embeddings(file, query, model_types, model_names, split_strategy, ch
             chunk_size,
             overlap_size,
             custom_separators.split(',') if custom_separators else None,
-            lang
+            lang,
+            custom_tokenizer_file
         )
 
         if use_custom_embedding:
@@ -324,7 +352,6 @@ def compare_embeddings(file, query, model_types, model_names, split_strategy, ch
             tokenizer, optimized_chunks = optimize_vocabulary(chunks)
             chunks = optimized_chunks
 
-
         results, search_time, vector_store = search_embeddings(
             chunks,
             embedding_model,
@@ -347,7 +374,10 @@ def compare_embeddings(file, query, model_types, model_names, split_strategy, ch
     results_df = pd.DataFrame(all_results)
     stats_df = pd.DataFrame(all_stats)
 
-    return results_df, stats_df
+    # Generate visualizations
+    fig = visualize_results(results_df, stats_df)
+
+    return results_df, stats_df, fig
 
 def format_results(results, stats):
     formatted_results = []
@@ -355,53 +385,22 @@ def format_results(results, stats):
         result = {
             "Model": stats["model"],
             "Content": doc.page_content,
+            "Embedding": doc.embedding if hasattr(doc, 'embedding') else None,
             **doc.metadata,
             **{k: v for k, v in stats.items() if k not in ["model"]}
         }
         formatted_results.append(result)
     return formatted_results
 
-import matplotlib.pyplot as plt
-import seaborn as sns
-from sklearn.manifold import TSNE
-
-def visualize_results(results_df, stats_df):
-    # Create a figure with subplots
-    fig, axs = plt.subplots(2, 2, figsize=(20, 20))
-    
-    # 1. Bar plot of search times
-    sns.barplot(x='model', y='search_time', data=stats_df, ax=axs[0, 0])
-    axs[0, 0].set_title('Search Time by Model')
-    axs[0, 0].set_xticklabels(axs[0, 0].get_xticklabels(), rotation=45, ha='right')
-    
-    # 2. Scatter plot of result diversity vs. rank correlation
-    sns.scatterplot(x='result_diversity', y='rank_correlation', hue='model', data=stats_df, ax=axs[0, 1])
-    axs[0, 1].set_title('Result Diversity vs. Rank Correlation')
-    
-    # 3. Box plot of content lengths
-    sns.boxplot(x='model', y='content_length', data=results_df, ax=axs[1, 0])
-    axs[1, 0].set_title('Distribution of Result Content Lengths')
-    axs[1, 0].set_xticklabels(axs[1, 0].get_xticklabels(), rotation=45, ha='right')
-    
-    # 4. t-SNE visualization of embeddings
-    embeddings = np.array(results_df['embedding'].tolist())
-    tsne = TSNE(n_components=2, random_state=42)
-    embeddings_2d = tsne.fit_transform(embeddings)
-    
-    sns.scatterplot(x=embeddings_2d[:, 0], y=embeddings_2d[:, 1], hue=results_df['model'], ax=axs[1, 1])
-    axs[1, 1].set_title('t-SNE Visualization of Result Embeddings')
-    
-    plt.tight_layout()
-    return fig
-
+# 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()) + ["Custom"], label="Embedding Model Types"),
-            gr.CheckboxGroup(choices=[model for models in MODELS.values() for model in models] + ["custom_model"], label="Embedding Models"),
+            gr.CheckboxGroup(choices=list(model_manager.list_models().keys()) + ["Custom"], label="Embedding Model Types"),
+            gr.CheckboxGroup(choices=[model for models in model_manager.list_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"),
@@ -412,7 +411,8 @@ def launch_interface(share=True):
             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")
+            gr.Slider(0, 1, step=0.1, value=0.3, label="Phonetic Matching Weight"),
+            gr.File(label="Custom Tokenizer File (Optional)")
         ],
         outputs=[
             gr.Dataframe(label="Results", interactive=False),
@@ -426,7 +426,21 @@ def launch_interface(share=True):
     tutorial_md = """
     # Advanced Embedding Comparison Tool Tutorial
 
-    ... (update the tutorial to include information about the new features) ...
+    This tool allows you to compare different embedding models and retrieval strategies for document search and similarity matching.
+
+    ## How to use:
+
+    1. Upload a file (optional) or use the default files in the system.
+    2. Enter a search query.
+    3. Select one or more embedding model types and specific models.
+    4. Choose a text splitting strategy and set chunk size and overlap.
+    5. Select a vector store type and search type.
+    6. Set the number of top results to retrieve.
+    7. Choose the language of your documents.
+    8. Optionally, use custom embeddings, optimize vocabulary, or adjust phonetic matching weight.
+    9. If you have a custom tokenizer, upload the file.
+
+    The tool will process your query and display results, statistics, and visualizations to help you compare the performance of different models and strategies.
     """
 
     iface = gr.TabbedInterface(