Update app.py

app.py

@@ -505,18 +505,64 @@ def optimize_vocabulary(texts, vocab_size=10000, min_frequency=2):
     
     return tokenizer, optimized_texts
     
-# New postprocessing function
-def rerank_results(results, query, reranker):
+import numpy as np
+from transformers import TextClassificationPipeline
+from typing import List, Union, Any
+
+def rerank_results(
+    results: List[Any], 
+    query: str, 
+    reranker: Union[TextClassificationPipeline, Any]
+) -> List[Any]:
+    """
+    Rerank search results using either a TextClassificationPipeline or a custom reranker.
+    
+    Args:
+        results: List of documents/results to rerank
+        query: Search query string
+        reranker: Either a HuggingFace TextClassificationPipeline or a custom reranker 
+                 with a rerank() method
+    
+    Returns:
+        List of reranked results
+    """
+    if not results:
+        return 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]
+        try:
+            pairs = [[query, doc.page_content] for doc in results]
+            
+            # Standard classification without specific function
+            predictions = reranker(pairs)
+            
+            # Extract scores, defaulting to 'score' key but falling back to other common keys
+            scores = []
+            for pred in predictions:
+                if isinstance(pred, dict):
+                    score = pred.get('score', 
+                           pred.get('probability',
+                           pred.get('confidence', 0.0)))
+                else:
+                    score = float(pred)
+                scores.append(score)
+                
+            # Sort in descending order (higher scores = better matches)
+            reranked_idx = np.argsort(scores)[::-1]
+            return [results[i] for i in reranked_idx]
+            
+        except Exception as e:
+            print(f"Warning: Reranking failed with error: {str(e)}")
+            return results
     else:
-        # For models with rerank method
-        return reranker.rerank(query, [doc.page_content for doc in results])
-
+        # For models with dedicated rerank method
+        try:
+            return reranker.rerank(query, [doc.page_content for doc in results])
+        except Exception as e:
+            print(f"Warning: Custom reranking failed with error: {str(e)}")
+            return 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):
     all_results = []
@@ -660,7 +706,9 @@ def automated_testing(file, query, test_params, expected_result=None):
             chunks = optimized_chunks
 
         if params['use_query_optimization']:
-            optimized_queries = optimize_query(query, params['query_optimization_model'])
+            optimized_queries = optimize_query(query, params['query_optimization_model'], params['chunks'] ,  params['embedding_model'] ,  params['vector_store_type'] , params['search_type'] ,  params['top_k'] )
+
+            #optimized_queries = optimize_query(query, )
             query = " ".join(optimized_queries)
 
         results, search_time, vector_store, results_raw = search_embeddings(