Upload 2 files

app.py

@@ -191,17 +191,35 @@ def create_paris_map(results_df):
 #         import traceback
 #         return f"Error: {str(e)}\n\n{traceback.format_exc()}", None
 
-# def search_restaurants(query, data_source, search_method, num_results, use_popularity):
-#     """Main search function that routes to appropriate search method"""
-#     if search_method == "Semantic Search" and use_semantic:
-#         return semantic_search(query, data_source, num_results, use_popularity)
-#     else:
-#         return keyword_search(query, data_source, num_results, use_popularity)
 
 def search_restaurants(query_input, data_source, num_results):
     n_candidates = 100
     query_clean = clean_text(query_input)
-    return get_recommendations(query_clean, n_candidates, num_results)
+    restaurant_ids = get_recommendations(query_clean, n_candidates, num_results, data_source)
+    
+    # Subset data for recommendedations
+    results = data[data["id"].isin(restaurant_ids)]
+    map_html = create_paris_map(results)
+
+    output = f"Found {len(results)} restaurants for '{query_input}'\n"
+    output += f"Data Source: {data_source}\n"
+
+    for idx, (_, row) in enumerate(results.iterrows(), 1):
+        name = row.get('name', 'Unknown')
+        rating = row.get('overall_rating', 'N/A')
+        reviews = row.get('review_count', 'N/A')
+        
+        output += f"{idx}. **{name}**\n"
+        output += f"   Rating: {rating} | Reviews: {reviews}\n"
+        output += "\n"
+        
+        if 'address' in row and pd.notna(row['address']):
+            addr = str(row['address'])[:100]
+            output += f"   Address: {addr}\n"
+        
+        output += "\n"
+
+    return output, map_html
 
 # Create Gradio interface
 with gr.Blocks(
@@ -225,7 +243,7 @@ with gr.Blocks(
         
         with gr.Column(scale=2):
             data_source = gr.Dropdown(
-                choices=["Michelin", "Google", "Yelp"],
+                choices=["Michelin Guide", "Google", "Yelp"],
                 value="Yelp",
                 label="Data Source",
                 info="Select restaurant data source"

main.py

@@ -4,6 +4,7 @@ import nltk
 import benepar
 import pandas as pd 
 import numpy as np
+from tqdm import tqdm
 from sklearn.metrics.pairwise import cosine_similarity 
 
 from utils.clean_text import clean_text 
@@ -30,7 +31,11 @@ with open("data/restaurant_by_source.json", "r") as f:
     restaurant_by_source = json.load(f)
 
 # Load precomputed TF-IDF features 
-restaurant_tfidf_features = np.load("data/toy_data_tfidf_features.npz") 
+# restaurant_tfidf_features = np.load("data/toy_data_tfidf_features.npz") 
+
+print("Computing TFIDF")
+tfidf_vectorizer = TFIDF_Vectorizer(load_vectorizer=False) 
+restaurant_tfidf_features = tfidf_vectorizer.compute_tfidf_matrix(data["review_text_clean"])
 
 # Extract embeddings
 data["embedding"] = data["embedding"].apply(
@@ -44,25 +49,28 @@ encoder = Encoder()
 # Initialize syntactic parser 
 parser = Parser() 
 
-# Initialize TF-IDF vectorizer 
-tfidf_vectorizer = TFIDF_Vectorizer(load_vectorizer=True) 
-
 def retrieve_candidates(query: str, n_candidates: int): 
+    print(f"Retrieving {n_candidates} candidates...")
+
     # Encode query 
+    print("[RETRIEVAL] Encoding query")
     query_emb = encoder.encode([query]).cpu().numpy() 
     
     # Semantic similarities 
+    print("[RETRIEVAL] Computing semantic similarities")
     desc_sem_sim = cosine_similarity(query_emb, all_desc_embeddings)[0] 
     
     # TF-IDF similarities 
+    print("[RETRIEVAL] Computing TF-IDF")
     tfidf_sim = tfidf_vectorizer.compute_tfidf_scores(query, restaurant_tfidf_features) 
     
     # Syntactic similarities 
+    print("[RETRIEVAL] Computing syntactic similarities")
     parsed_query = parser.parse_text(query)
     parsed_query = parser.subtree_set(parsed_query)
 
     syn_sims = []
-    for trees_list in data["syntactic_tree"]:
+    for trees_list in tqdm(data["syntactic_tree"], total=len(data), desc="[RETRIEVAL] Computing syntactic similarities"):
         review_sims = []
         for review_tree_subs in trees_list:
             if review_tree_subs is None:
@@ -78,10 +86,13 @@ def retrieve_candidates(query: str, n_candidates: int):
     # Get top N candidates for Stage 2 reranking
     candidates_idx = np.argsort(combined_stage1_scores)[-n_candidates:][::-1]
 
+    print(f"[RETRIEVAL] Results: {candidates_idx}")
+
     return candidates_idx
 
 
 def rerank(candidates_idx: np.ndarray, n_rec: int = 10, data_source: str = None) -> list:
+    print("Reranking...")
     
     # Get popularity scores for stage 1 candidates 
     rerank_scores = data.loc[candidates_idx, "pop_score"].values
@@ -93,6 +104,12 @@ def rerank(candidates_idx: np.ndarray, n_rec: int = 10, data_source: str = None)
     # Get restaurant_id for final recommendations 
     restaurant_ids = data.loc[topN_reranked_global_idx, "id"].tolist() 
     
+    # Filter to only data_source
+    print(f"[RERANK] Filtering to only source - {data_source}")
+    restaurant_by_source_set = set(restaurant_by_source[data_source])
+    restaurant_ids = [x for x in restaurant_ids if x in restaurant_by_source_set]
+    
+    print(f"[RERANK] Final recommendations: {restaurant_ids}")
     return restaurant_ids
 
 def get_recommendations(query: str, n_candidates: int = 100, n_rec: int = 30, data_source: str = None):