Update app.py

app.py

@@ -49,6 +49,8 @@ def download_nltk_resources():
 
 download_nltk_resources()
 
+nltk.download('punkt')
+
 FILES_DIR = './files'
 
 # Model Management
@@ -159,9 +161,9 @@ def phonetic_match(text, query, method='levenshtein_distance', apply_phonetic=Tr
 
 def optimize_query(query, llm_model):
     llm = HuggingFacePipeline.from_model_id(
-        model_id=llm_model,
+        model_id="google/flan-t5-large",
         task="text2text-generation",
-        model_kwargs={"temperature": 0, "max_length": 64},
+        model_kwargs={"do_sample": True, "temperature": 0.7, "max_new_tokens": 512},
     )
     multi_query_retriever = MultiQueryRetriever.from_llm(
         retriever=get_retriever(vector_store, search_type, search_kwargs),
@@ -391,6 +393,8 @@ def visualize_results(results_df, stats_df):
     
     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')
     
     sns.scatterplot(x='result_diversity', y='rank_correlation', hue='model', data=stats_df, ax=axs[0, 1])
@@ -398,6 +402,7 @@ def visualize_results(results_df, stats_df):
     
     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')
     
     embeddings = np.array([embedding for embedding in results_df['embedding'] if isinstance(embedding, np.ndarray)])
@@ -514,6 +519,8 @@ def compare_embeddings(file, query, embedding_models, custom_embedding_model, sp
 
         stats = calculate_statistics(results_raw, search_time, vector_store, num_tokens, embedding_model, query, top_k, expected_result)
         stats["model"] = f"{model_type} - {model_name}"
+        stats["model_type"] = model_type
+        stats["model_name"] = model_name
         stats.update(settings)
 
         formatted_results = format_results(results_raw, stats)
@@ -605,6 +612,8 @@ def automated_testing(file, query, test_params, expected_result=None):
 
         stats = calculate_statistics(results_raw, search_time, vector_store, num_tokens, embedding_model, query, params['top_k'], expected_result)
         stats["model"] = f"{params['model_type']} - {params['model_name']}"
+        stats["model_type"] = model_type
+        stats["model_name"] = model_name
         stats.update(params)
 
         all_results.extend(format_results(results_raw, stats))
@@ -705,7 +714,7 @@ Provide your suggestions in a Python dictionary format."""
     llm = HuggingFacePipeline.from_model_id(
         model_id="google/flan-t5-large",
         task="text2text-generation",
-        model_kwargs={"temperature": 0.7, "max_length": 512},  # Changed max_length to max_new_tokens
+        model_kwargs={"do_sample": True, "temperature": 0.7, "max_new_tokens": 512},
     )
 
     # Generate suggestions