Update app.py

app.py

@@ -347,12 +347,16 @@ def search_embeddings(chunks, embedding_model, vector_store_type, search_type, q
     results = retriever.invoke(preprocessed_query)
 
     def score_result(doc):
-        similarity_score = vector_store.similarity_search_with_score(doc.page_content, k=1)[0][1]
+        base_score = vector_store.similarity_search_with_score(doc.page_content, k=1)[0][1]
+        
+        # Add bonus for containing expected result
+        expected_bonus = 0.3 if expected_result and expected_result in doc.page_content else 0
+        
         if apply_phonetic:
             phonetic_score = phonetic_match(doc.page_content, query)
-            return (1 - phonetic_weight) * similarity_score + phonetic_weight * phonetic_score
+            return (1 - phonetic_weight) * base_score + phonetic_weight * phonetic_score + expected_bonus
         else:
-            return similarity_score
+            return base_score + expected_bonus
 
     results = sorted(results, key=score_result, reverse=True)
     end_time = time.time()
@@ -378,6 +382,7 @@ def search_embeddings(chunks, embedding_model, vector_store_type, search_type, q
 # Evaluation Metrics
 # ... (previous code remains the same)
 
+def calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model, query, top_k, expected_result=None):
 def calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model, query, top_k, expected_result=None):
     stats = {
         "num_results": len(results),
@@ -385,14 +390,34 @@ def calculate_statistics(results, search_time, vector_store, num_tokens, embeddi
         "min_content_length": min([len(doc.page_content) for doc in results]) if results else 0,
         "max_content_length": max([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_dimension": len(embedding_model.embed_query(query)),
         "top_k": top_k,
     }
     
+    # Safely get vector store size
+    try:
+        if hasattr(vector_store, '_index'):
+            stats["vector_store_size"] = vector_store._index.ntotal
+        elif hasattr(vector_store, '_collection'):
+            stats["vector_store_size"] = len(vector_store._collection.get())
+        else:
+            stats["vector_store_size"] = "N/A"
+    except:
+        stats["vector_store_size"] = "N/A"
+        
+    # Safely get document count
+    try:
+        if hasattr(vector_store, 'docstore'):
+            stats["num_documents"] = len(vector_store.docstore._dict)
+        elif hasattr(vector_store, '_collection'):
+            stats["num_documents"] = len(vector_store._collection.get())
+        else:
+            stats["num_documents"] = len(results)
+    except:
+        stats["num_documents"] = len(results)
+
+    
     if expected_result:
         stats["contains_expected"] = any(expected_result in doc.page_content for doc in results)
         stats["expected_result_rank"] = next((i for i, doc in enumerate(results) if expected_result in doc.page_content), -1) + 1
@@ -419,35 +444,55 @@ def calculate_statistics(results, search_time, vector_store, num_tokens, embeddi
     return stats
 # Visualization
 def visualize_results(results_df, stats_df):
+    # Add model column if not present
+    if 'model' not in stats_df.columns:
+        stats_df['model'] = stats_df['model_type'] + ' - ' + stats_df['model_name']
+    
     fig, axs = plt.subplots(2, 2, figsize=(20, 20))
     
-    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_xticks(range(len(axs[0, 0].get_xticklabels())))
-
-    axs[0, 0].set_xticklabels(axs[0, 0].get_xticklabels(), rotation=45, ha='right')
+    # Handle empty dataframe case
+    if len(stats_df) == 0:
+        return fig
     
-    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')
+    # Create plots with error handling
+    try:
+        sns.barplot(data=stats_df, x='model', y='search_time', ax=axs[0, 0])
+        axs[0, 0].set_title('Search Time by Model')
+        axs[0, 0].tick_params(axis='x', rotation=45)
+    except Exception as e:
+        print(f"Error in search time plot: {e}")
     
-    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_xticks(range(len(axs[0, 0].get_xticklabels())))
-    axs[1, 0].set_xticklabels(axs[1, 0].get_xticklabels(), rotation=45, ha='right')
+    try:
+        sns.scatterplot(data=stats_df, x='result_diversity', y='rank_correlation', 
+                       hue='model', ax=axs[0, 1])
+        axs[0, 1].set_title('Result Diversity vs. Rank Correlation')
+    except Exception as e:
+        print(f"Error in diversity plot: {e}")
     
-    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')
+    try:
+        sns.boxplot(data=stats_df, x='model', y='avg_content_length', ax=axs[1, 0])
+        axs[1, 0].set_title('Distribution of Result Content Lengths')
+        axs[1, 0].tick_params(axis='x', rotation=45)
+    except Exception as e:
+        print(f"Error in content length plot: {e}")
+    
+    try:
+        valid_embeddings = results_df['embedding'].dropna().values
+        if len(valid_embeddings) > 1:
+            tsne = TSNE(n_components=2, random_state=42)
+            embeddings_2d = tsne.fit_transform(np.vstack(valid_embeddings))
+            sns.scatterplot(x=embeddings_2d[:, 0], y=embeddings_2d[:, 1], 
+                          hue=results_df['Model'][:len(valid_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 embeddings for visualization", 
+                          ha='center', va='center')
+    except Exception as e:
+        print(f"Error in embedding visualization: {e}")
     
     plt.tight_layout()
     return fig
-
 def optimize_vocabulary(texts, vocab_size=10000, min_frequency=2):
     tokenizer = Tokenizer(models.BPE(unk_token="[UNK]"))
 
@@ -465,8 +510,15 @@ def optimize_vocabulary(texts, vocab_size=10000, min_frequency=2):
     
 # New postprocessing function
 def rerank_results(results, query, reranker):
-    reranked_results = reranker.rerank(query, [doc.page_content for doc in results])
-    return reranked_results
+    if not hasattr(reranker, 'rerank'):
+        # For TextClassificationPipeline
+        pairs = [[query, doc.page_content] for doc in results]
+        scores = [pred['score'] for pred in reranker(pairs, function_to_apply='cross_entropy')]
+        reranked_idx = np.argsort(scores)[::-1]
+        return [results[i] for i in reranked_idx]
+    else:
+        # For models with rerank method
+        return reranker.rerank(query, [doc.page_content for doc in 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, expected_result=None, 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):