Update app.py

app.py

@@ -185,16 +185,35 @@ def search_embeddings(chunks, embedding_model, vector_store_type, search_type, q
 
     return results[:top_k], end_time - start_time, vector_store
 
-def calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model):
-    return {
+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,
         "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),
         "num_tokens": num_tokens,
-        "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"
+        "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",
+        "embedding_dimension": len(embedding_model.embed_query(query)),
+        "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)
+        stats["result_diversity"] = 1 - np.mean(pairwise_similarities[np.triu_indices(len(embeddings), k=1)])
+    else:
+        stats["result_diversity"] = "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]
+    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
@@ -236,7 +255,7 @@ def optimize_vocabulary(texts, vocab_size=10000, min_frequency=2):
     
     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):
+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):
     all_results = []
     all_stats = []
     settings = {
@@ -273,6 +292,7 @@ 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,
@@ -284,7 +304,7 @@ def compare_embeddings(file, query, model_types, model_names, split_strategy, ch
             phonetic_weight
         )
 
-        stats = calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model)
+        stats = calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model, query, top_k)
         stats["model"] = f"{model_type} - {model_name}"
         stats.update(settings)
 
@@ -309,6 +329,39 @@ def format_results(results, stats):
         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
+
 def launch_interface(share=True):
     iface = gr.Interface(
         fn=compare_embeddings,
@@ -331,7 +384,8 @@ def launch_interface(share=True):
         ],
         outputs=[
             gr.Dataframe(label="Results", interactive=False),
-            gr.Dataframe(label="Statistics", 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"