Update app.py

app.py

@@ -4,10 +4,15 @@ import numpy as np
 from ast import literal_eval
 import pandas as pd
 
+import asyncio
 from gradio_folium import Folium
 import folium
 from huggingface_hub import InferenceClient
 from geopy.geocoders import Nominatim
+from collections import OrderedDict
+from geopy.adapters import AioHTTPAdapter
+import nest_asyncio
+nest_asyncio.apply()
 
 from examples import (
     description_sf,
@@ -17,15 +22,6 @@ from examples import (
     df_examples
 )
 
-geolocator = Nominatim(user_agent="HF-trip-planner")
-
-def get_coordinates(address):
-    location = geolocator.geocode(address)
-    if location:
-        return (location.latitude, location.longitude)
-    else:
-        return None
-
 repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 llm_client = InferenceClient(model=repo_id, timeout=180)
 
@@ -60,15 +56,61 @@ def parse_llm_output(output):
     return dataframe, rationale
 
 
-def get_coordinates_row(row):
-    coords = get_coordinates(row["name"])
-    if coords is not None:
-        row["lat"], row["lon"] = coords
-    return row
+class AsyncLRUCache:
+    def __init__(self, maxsize=100):
+        self.cache = OrderedDict()
+        self.maxsize = maxsize
+
+    async def get(self, key):
+        if key not in self.cache:
+            return None
+        self.cache.move_to_end(key)
+        return self.cache[key]
+
+    async def set(self, key, value):
+        if key in self.cache:
+            self.cache.move_to_end(key)
+        self.cache[key] = value
+        if len(self.cache) > self.maxsize:
+            self.cache.popitem(last=False)
+
+
+# Instantiate the cache
+cache = AsyncLRUCache(maxsize=500)
+
+async def geocode_address(address):
+    # Check if the result is in cache
+    cached_location = await cache.get(address)
+    if cached_location:
+        return cached_location
+
+    # If not in cache, perform the geolocation request
+    async with Nominatim(
+        user_agent="HF-trip-planner",
+        adapter_factory=AioHTTPAdapter,
+    ) as geolocator:
+        location = await geolocator.geocode(address, timeout=10)
+        if location:
+            # Save the result in cache for future use
+            await cache.set(address, location)
+        return location
+    
+async def ageocode_addresses(addresses):
+    tasks = [geocode_address(address) for address in addresses]
+    locations = await asyncio.gather(*tasks)
+    return locations
+
+def geocode_addresses(addresses):
+    loop = asyncio.get_event_loop()
+    result = loop.run_until_complete(ageocode_addresses(addresses))
+    return result
 
 
 def create_map_from_markers(dataframe):
-    dataframe = dataframe.apply(get_coordinates_row, axis=1)
+    locations = geocode_addresses(dataframe["name"])
+    dataframe["lat"] = [location.latitude if location else None for location in locations]
+    dataframe["lon"] = [location.longitude if location else None for location in locations]
+
     f_map = Map(
         location=[dataframe["lat"].mean(), dataframe["lon"].mean()],
         zoom_start=5,
@@ -106,7 +148,8 @@ def run_display(text):
 def select_example(df, data: gr.SelectData):
     row = df.iloc[data.index[0], :]
     dataframe, rationale = parse_llm_output(row["output"])
-    return row["description"], create_map_from_markers(dataframe), rationale
+    map = create_map_from_markers(dataframe)
+    return row["description"], map, rationale
 
 
 with gr.Blocks(
@@ -127,7 +170,7 @@ with gr.Blocks(
     example_dataframe, example_rationale = parse_llm_output(output_example_sf)
     display_rationale = gr.Markdown(example_rationale)
     starting_map = create_map_from_markers(example_dataframe)
-    map = Folium(value=starting_map, height=700, label="Chosen locations")
+    map = Folium(value=starting_map, height=600, label="Chosen locations")
     button.click(run_display, inputs=[text], outputs=[map, display_rationale])
 
     gr.Markdown("### Other examples")