first draft, with the gcp bucket

.gitignore

@@ -182,4 +182,8 @@ gradio_cached_examples/
 models/__pycache__/
 setup/__pycache__/
 setup/gcp_default_creds.json
-.config/*.json
+.config/*.json
+gradio_cached_examples/
+src/gradio_cached_examples/
+database/
+european-city-data/

requirements.txt

@@ -1,10 +1,12 @@
 sentence-transformers
 gradio
 gradio_client
-pymongo==4.6.2
 python-dotenv
 google-cloud-aiplatform
 google-cloud
-vertexai==1.43.0
+vertexai==1.64.0
 huggingface_hub
-certifi==2021.5.30
+certifi==2024.7.4
+python-dotenv==1.0.1
+bitsandbytes
+accelerate

src/__init__.py

@@ -0,0 +1,11 @@
+from src.vectordb.vectordb import *
+from src.vectordb.helpers import *
+from src.vectordb.lancedb_init import *
+
+from src.sustainability.s_fairness import *
+from src.information_retrieval.info_retrieval import *
+from src.augmentation.prompt_generation import *
+from src.text_generation.model_init import *
+from src.text_generation.text_generation import *
+
+from src.data_directories import *

src/app.py

@@ -0,0 +1,78 @@
+from typing import Optional
+import gradio as gr
+import os, sys
+sys.path.append("./src")
+print(os.getcwd())
+from src.pipeline import pipeline
+
+
+def clear():
+    return None, None, None
+
+
+def generate_text(query_text, model_name: Optional[str], is_sustainable: Optional[bool], tokens: Optional[int] = 1024,
+                  temp: Optional[float] = 0.49):
+    if is_sustainable:
+        sustainability = 1
+    else:
+        sustainability = 0
+    pipeline_response = pipeline(
+        query=query_text,
+        model_name=model_name,
+        sustainability= sustainability
+    )
+    return pipeline_response
+
+
+examples = [["I'm planning a vacation to France. Can you suggest a one-week itinerary including must-visit places and "
+             "local cuisines to try?", "GPT-4"],
+            ["I want to explore off-the-beaten-path destinations in Europe, any suggestions?", "Gemini-1.0-pro"],
+            ["Suggest some cities that can be visited from London and are very rich in history and culture.",
+             "Gemini-1.0-pro"],
+            ]
+
+with gr.Blocks() as demo:
+    gr.HTML("""<center><h1 style='font-size:xx-large;'>🇪🇺 Euro City Recommender using Gemini & Gemma 🇪🇺</h1><br><h3>Gemini 
+    & Gemma Sprints 2024 submissions by Ashmi Banerjee. </h3></center> <br><p>We're testing the compatibility of 
+    Retrieval Augmented Generation (RAG) implementations with Google's <b>Gemma-2b-it</b> & <b>Gemini 1.0 Pro</b> 
+    models through HuggingFace and VertexAI, respectively, to generate travel recommendations. This early version (read 
+    quick and dirty implementation) aims to see if functionalities work smoothly. It relies on Wikipedia abstracts 
+    from 160 European cities to provide answers to your questions. Please be kind with it, as it's a work in progress! 
+    </p> <br>Google Cloud credits are provided for this project. </p>
+    """)
+
+    with gr.Group():
+        query = gr.Textbox(label="Query", placeholder="Ask for your city recommendation here!")
+        sustainable = gr.Checkbox(label="Sustainable", info="If you want sustainable recommendations for "
+                                                                        "hidden gems?")
+        model = gr.Dropdown(
+            ["GPT-4", "Gemini-1.0-pro"], label="Model", info="Select your model. Will add more "
+                                                                                  "models "                          
+                                                                                  "later!",
+        )
+        output = gr.Textbox(label="Generated Results", lines=4)
+
+        with gr.Accordion("Settings", open=False):
+            max_new_tokens = gr.Slider(label="Max new tokens", value=1024, minimum=0, maximum=8192, step=64,
+                                       interactive=True,
+                                       visible=True, info="The maximum number of output tokens")
+            temperature = gr.Slider(label="Temperature", step=0.01, minimum=0.01, maximum=1.0, value=0.49,
+                                    interactive=True,
+                                    visible=True, info="The value used to module the logits distribution")
+    with gr.Group():
+        with gr.Row():
+            submit_btn = gr.Button("Submit", variant="primary")
+            clear_btn = gr.Button("Clear", variant="secondary")
+            cancel_btn = gr.Button("Cancel", variant="stop")
+    submit_btn.click(generate_text, inputs=[query, model, sustainable], outputs=[output])
+    clear_btn.click(clear, inputs=[], outputs=[query, model, output])
+    cancel_btn.click(clear, inputs=[], outputs=[query, model, output])
+
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples, inputs=[query, model], label="Examples", fn=generate_text, outputs=[output],
+        cache_examples=True,
+    )
+
+if __name__ == "__main__":
+    demo.launch(show_api=False)

src/augmentation/prompt_generation.py

@@ -0,0 +1,185 @@
+from information_retrieval import info_retrieval as ir
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(encoding='utf-8', level=logging.DEBUG)
+
+
+def generate_prompt(query, context, template=None):
+    """
+    Function that generates the prompt given the user query and retrieved context. A specific prompt template will be
+    used if provided, otherwise the default base_prompt template is used.
+
+    Args: 
+        - query: str
+        - context: list[dict]
+        - template: str
+    
+    """
+
+    if template:
+        SYS_PROMPT = template
+    else:
+        SYS_PROMPT = """You are an AI recommendation system. Your task is to recommend cities in Europe for travel 
+        based on the user's question. You should use the provided contexts to suggest a list of the 3 best cities 
+        that are best suited to the user's question, as well as the month of travel. If the user has already provided 
+        the month of travel in the question, use the same month; otherwise, provide the ideal month of travel. Each 
+        recommendation should also contain an explanation of why it is being recommended, based on the context. Your 
+        answer must begin with "I recommend " followed by the city name and why you recommended it. Your answers are 
+        correct, high-quality, and written by a domain expert. If the provided context does not contain the answer, 
+        simply state, "The provided context does not have the answer." """
+
+    USER_PROMPT = """ Question: {} Which city do you recommend and why?
+
+    Context: Here are the options: {} 
+    
+    Answer:
+
+    """
+
+    formatted_prompt = f"{USER_PROMPT.format(query, context)}"
+    messages = [
+        {"role": "system", "content": SYS_PROMPT},
+        {"role": "user", "content": formatted_prompt}
+    ]
+
+    return messages
+
+
+def format_context(context):
+    """
+    Function that formats the retrieved context in a way that is easy for the LLM to understand. 
+
+    Args:
+        - context: list[dict]; retrieved context 
+    
+    """
+
+    formatted_context = ''
+
+    for i, (city, info) in enumerate(context.items()):
+
+        text = f"Option {i + 1}: {city} is a city in {info['country']}."
+        info_text = f"Here is some information about the city. {info['text']}"
+
+        attractions_text = "Here are some attractions: "
+        att_flag = 0
+        restaurants_text = "Here are some places to eat/drink: "
+        rest_flag = 0
+
+        hotels_text = "Here are some hotels: "
+        hotel_flag = 0
+
+        if len(info['listings']):
+            for listing in info['listings']:
+                if listing['type'] in ['see', 'do', 'go', 'view']:
+                    att_flag = 1
+                    attractions_text += f"{listing['name']} ({listing['description']}), "
+                elif listing['type'] in ['eat', 'drink']:
+                    rest_flag = 1
+                    restaurants_text += f"{listing['name']} ({listing['description']}), "
+                else:
+                    hotel_flag = 1
+                    hotels_text += f"{listing['name']} ({listing['description']}), "
+
+        # If we add sustainability in the end then it could get truncated because of context window
+        if "sustainability" in info:
+            if info['sustainability']['month'] == 'No data available':
+                sfairness_text = "This city has no sustainability (or s-fairness) score available."
+
+            else:
+                sfairness_text = f"The sustainability (or s-fairness) score for {city} in {info['sustainability']['month']} is {info['sustainability']['s-fairness']}. \n "
+
+            text += sfairness_text
+
+        text += info_text
+
+        if att_flag:
+            text += f"\n{attractions_text}"
+
+        if rest_flag:
+            text += f"\n{restaurants_text}"
+
+        if hotel_flag:
+            text += f"\n{hotels_text}"
+
+        formatted_context += text + "\n\n "
+
+    return formatted_context
+
+
+def augment_prompt(query, context, **params):
+    """
+    Function that accepts the user query as input, obtains relevant documents and augments the prompt with the
+    retrieved context, which can be passed to the LLM.
+
+    Args: - query: str - context: retrieved context, must be formatted otherwise the LLM cannot understand the nested
+    dictionaries! - sustainability: bool; if true, then the prompt is appended to instruct the LLM to use s-fairness
+    scores while generating the answer - params: key-value parameters to be passed to the get_context function; sets
+    the limit of results and whether to rerank the results
+    
+    """
+
+    # what about the cities without s-fairness scores? i.e. they don't have seasonality data 
+
+    prompt_with_sustainability = """You are an AI recommendation system. Your task is to recommend cities in Europe 
+    for travel based on the user's question. You should use the provided contexts to suggest the city that is best 
+    suited to the user's question. You recommend a list of the top 3 most sustainable cities to the user, as well as 
+    the best month of travel. Each recommendation should also contain an explanation of why it is being recommended, 
+    on sustainability grounds based on the context. The context contains a sustainability score for each city, 
+    also known as the s-fairness score, along with the ideal month of travel. A lower s-fairness score indicates that 
+    the city is a better destination for the month provided. A city without a sustainability score should not be 
+    considered. You should only consider the s-fairness score while choosing the best city. However, your answer 
+    should not contain the numeric score itself or any mention of the sustainability score. Your answer must begin 
+    with "I recommend " followed by the city names and why you recommended it. Your answers are correct, high-quality, 
+    and written by a domain expert. If the provided context does not contain the answer, simply state, "The provided 
+    context does not have the answer. """
+
+    # format context
+    formatted_context = format_context(context)
+
+    if "sustainability" in params["params"] and params["params"]["sustainability"]:
+        prompt = generate_prompt(query, formatted_context, prompt_with_sustainability)
+    else:
+        prompt = generate_prompt(query, formatted_context)
+
+    return prompt
+
+
+def test():
+    context_params = {
+        'limit': 3,
+        'reranking': 0, 
+        'sustainability': 0
+    }
+
+    query = "Suggest some places to visit during winter. I like hiking, nature and the mountains and I enjoy skiing " \
+            "in winter. "
+
+    # without sustainability
+    context = ir.get_context(query, **context_params)
+    # formatted_context = format_context(context)
+
+    without_sfairness = augment_prompt(
+        query=query,
+        context=context,
+        params=context_params
+    )
+
+    # with sustainability
+    context_params.update({'sustainability': 1})
+    s_context = ir.get_context(query, **context_params)
+    # s_formatted_context = format_context(s_context)
+
+    with_sfairness = augment_prompt(
+        query=query,
+        context=s_context,
+        params=context_params
+    )
+
+    return with_sfairness
+
+
+if __name__ == "__main__":
+    prompt = test()
+    print(prompt)

src/data_directories.py

@@ -0,0 +1,15 @@
+import os
+
+current = os.path.dirname(os.path.realpath(''))
+parent = os.path.dirname(current)
+
+data_parent_dir = "../../european-city-data/"
+data_dir = data_parent_dir + "data-sources/"
+wikivoyage_docs_dir = data_dir + "wikivoyage/"
+wikivoyage_listings_dir = wikivoyage_docs_dir + "listings/"
+database_dir = "../../database/wikivoyage/"
+seasonality_dir = data_parent_dir + "computed/seasonality/"
+popularity_dir = data_parent_dir + "computed/popularity/"
+cities_csv = data_parent_dir + "city_abstracts_embeddings.csv"
+prompts_dir = data_parent_dir + "rag-sustainability/prompts/"
+results_dir = data_parent_dir + "rag-sustainability/results/"

src/information_retrieval/info_retrieval.py

@@ -0,0 +1,267 @@
+import sys
+import re
+import os
+import json
+sys.path.append("../")
+from src.vectordb import vectordb
+from src.sustainability import s_fairness
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(encoding='utf-8', level=logging.DEBUG)
+
+
+def get_travel_months(query):
+    """
+
+    Function to parse the user's query and search if month of travel has been provided by the user.
+
+    Args:
+    - query: str
+
+    """
+    months = [
+        "January", "February", "March", "April", "May", "June",
+        "July", "August", "September", "October", "November", "December"
+    ]
+
+    seasons = {
+        "spring": ["March", "April", "May"],
+        "summer": ["June", "July", "August"],
+        "fall": ["September", "October", "November"],
+        "autumn": ["September", "October", "November"],
+        "winter": ["December", "January", "February"]
+    }
+
+    months_in_query = []
+
+    for month in months:
+        if re.search(r'\b' + month + r'\b', query, re.IGNORECASE):
+            months_in_query.append(month)
+
+    # Check for seasons in the query
+    for season, season_months in seasons.items():
+        if re.search(r'\b' + season + r'\b', query, re.IGNORECASE):
+            months_in_query += season_months
+
+    # Return None if neither months nor seasons are found
+    return months_in_query
+
+
+def get_wikivoyage_context(query, limit=10, reranking=0):
+    """
+
+    Function to retrieve the relevant documents and listings from the wikivoyage database. Works in two steps:
+    (i) the relevant cities are returned by the wikivoyage_docs table and (ii) then passed on to the wikivoyage listings database to retrieve further information.
+    The user can pass a limit of how many results the search should return as well as whether to perform reranking (uses a CrossEncoderReranker)
+
+    Args:
+        - query: str
+        - limit: int
+        - reranking: bool
+
+    """
+
+    # limit = params['limit']
+    # reranking = params['reranking']
+
+    docs = vectordb.search_wikivoyage_docs(query, limit, reranking)
+    logger.info("Finished getting chunked wikivoyage docs.")
+
+    results = {}
+    for doc in docs:
+        results[doc['city']] = {key: value for key, value in doc.items() if key != 'city'}
+        results[doc['city']]['listings'] = []
+
+    cities = [result['city'] for result in docs]
+
+    listings = vectordb.search_wikivoyage_listings(query, cities, limit, reranking)
+    logger.info("Finished getting wikivoyage listings.")
+    # logger.info(type(docs), type(listings))
+
+    for listing in listings:
+        # logger.info(listing['city'])
+        results[listing['city']]['listings'].append({
+            'type': listing['type'],
+            'name': listing['title'],
+            'description': listing['description']
+        })
+
+    logger.info("Returning retrieval results.")
+    return results
+
+
+def get_sustainability_scores(query, destinations):
+    """
+
+    Function to get the s-fairness scores for each destination for the given month (or the ideal month of travel if the user hasn't provided a month).
+    If multiple months are provided (or season), then the month with the minimum s-fairness score is chosen for the city.
+
+    Args:
+        - query: str
+        - destinations: list
+
+    """
+
+    result = []  # list of dicts of the format {city: <city>, month: <month>, }
+    city_scores = {}
+
+    months = get_travel_months(query)
+    logger.info("Finished parsing query for months.")
+
+    for city in destinations:
+        if city not in city_scores:
+            city_scores[city] = []
+
+        if not months:  # no month(s) or seasons provided by the user
+            city_scores[city].append(s_fairness.compute_sfairness_score(city))
+        else:
+            for month in months:
+                city_scores[city].append(s_fairness.compute_sfairness_score(city, month))
+
+    logger.info("Finished getting s-fairness scores.")
+
+    for city, scores in city_scores.items():
+
+        no_result = 0
+        for score in scores:
+            if not score['month']:
+                no_result = 1
+                result.append({
+                    'city': city,
+                    'month': 'No data available',
+                    's-fairness': 'No data available'
+                })
+                break
+
+        if not no_result:
+            min_score = min(scores, key=lambda x: x['s-fairness'])
+            result.append({
+                'city': city,
+                'month': min_score['month'],
+                's-fairness': min_score['s-fairness']
+            })
+
+    logger.info("Returning s-fairness results.")
+    return result
+
+
+def get_cities(context):
+    """
+    Only to be used for testing! Function that returns a list of cities with their s-fairness scores, provided the retrieved context
+
+    Args:
+        - context: dict
+
+    """
+
+    recommended_cities = []
+
+    for city, info in context.items():
+        city_info = {
+            'city': city,
+            'country': info['country']
+        }
+
+        if "sustainability" in info:
+            city_info['month'] = info['sustainability']['month']
+            city_info['s-fairness'] = info['sustainability']['s-fairness']
+
+        recommended_cities.append(city_info)
+
+    if "sustainability" in info:
+        def get_s_fairness_value(item):
+            s_fairness = item['s-fairness']
+            if s_fairness == 'No data available':
+                return float('inf')  # Assign a high value for "No data available"
+            return s_fairness
+
+        # Sort the list using the custom key
+        sorted_cities = sorted(recommended_cities, key=get_s_fairness_value)
+        return sorted_cities
+
+    else:
+        return recommended_cities
+
+
+def get_context(query, **params):
+    """
+
+    Function that returns all the context: from the database, as well as the respective s-fairness scores for the
+    destinations. The default does not consider S-Fairness scores, i.e. to append sustainability scores, a non-zero
+    parameter "sustainability" needs to be explicitly passed to params.
+
+    Args:
+        - query: str
+        - params: dict; contains value of the limit and reranking (and sustainability)
+
+    """
+
+    limit = 3
+    reranking = 1
+
+    if 'limit' in params:
+        limit = params['limit']
+
+    if 'reranking' in params:
+        reranking = params['reranking']
+
+    wikivoyage_context = get_wikivoyage_context(query, limit, reranking)
+    recommended_cities = wikivoyage_context.keys()
+
+    if 'sustainability' in params and params['sustainability']:
+        s_fairness_scores = get_sustainability_scores(query, recommended_cities)
+
+        for score in s_fairness_scores:
+            wikivoyage_context[score['city']]['sustainability'] = {
+                'month': score['month'],
+                's-fairness': score['s-fairness']
+            }
+
+    return wikivoyage_context
+
+
+def test():
+    queries = []
+    query = "Suggest some places to visit during winter. I like hiking, nature and the mountains and I enjoy skiing " \
+            "in winter. "
+
+    context = None
+
+    try:
+        context = get_context(query, sustainability=1)
+        # cities = get_cities(context)
+        # print(cities)
+    except FileNotFoundError as e:
+        try:
+            vectordb.create_wikivoyage_docs_db_and_add_data()
+            vectordb.create_wikivoyage_listings_db_and_add_data()
+
+            try:
+                context = get_context(query, sustainability=1)
+                # cities = get_cities(context)
+                # print(cities)
+            except Exception as e:
+                exc_type, exc_obj, exc_tb = sys.exc_info()
+                fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+                logger.error(f"Error while getting context: {e}, {(exc_type, fname, exc_tb.tb_lineno)}")
+
+        except Exception as e:
+            logger.error(f"Error while creating DB: {e}")
+
+    except Exception as e:
+        exc_type, exc_obj, exc_tb = sys.exc_info()
+        fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+        logger.error(f"Error while getting context: {e}, {(exc_type, fname, exc_tb.tb_lineno)}")
+
+    file_path = os.path.join(os.getcwd(), "test_results", "test_result.json")
+    with open(file_path, 'w') as file:
+        json.dump(context, file)
+
+    return context
+
+
+if __name__ == "__main__":
+    context = test()
+
+    print(context)

src/pipeline.py

@@ -0,0 +1,128 @@
+"""
+Main file to execute the TRS Pipeline.
+"""
+import sys
+from augmentation import prompt_generation as pg
+from information_retrieval import info_retrieval as ir
+from text_generation.models import (
+    Llama3,
+    Mistral,
+    Gemma2,
+    Llama3Point1,
+    Llama3Instruct,
+    MistralInstruct,
+    Llama3Point1Instruct,
+    Phi3SmallInstruct,
+    GPT4,
+    Gemini,
+)
+from text_generation import text_generation as tg
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(encoding='utf-8', level=logging.DEBUG)
+
+TEST_DIR = "../tests/"
+MODELS = {
+    'GPT-4': GPT4,
+    'Llama3': Llama3,
+    'Mistral': Mistral,
+    'Gemma2': Gemma2,
+    'Llama3.1': Llama3Point1,
+    'Llama3-Instruct': Llama3Instruct,
+    'Mistral-Instruct': MistralInstruct,
+    'Llama3.1-Instruct': Llama3Point1Instruct,
+    'Phi3-Instruct': Phi3SmallInstruct,
+    'Gemini-1.0-pro': Gemini,
+}
+
+
+def pipeline(query: str, model_name: str, test: int = 0, **params):
+    """
+    
+    Executes the entire RAG pipeline, provided the query and model class name.
+
+    Args: 
+        - query: str
+        - model_name: string, one of the following: Llama3, Mistral, Gemma2, Llama3Point1
+        - test: whether the pipeline is running a test
+        - params: 
+            - limit (number of results to be retained) 
+            - reranking (binary, whether to rerank results using ColBERT or not)
+            - sustainability
+
+    
+    """
+
+    model = MODELS[model_name]
+
+    context_params = {
+        'limit': 5,
+        'reranking': 0,
+        'sustainability': 0,
+    }
+
+    if 'limit' in params:
+        context_params['limit'] = params['limit']
+
+    if 'reranking' in params:
+        context_params['reranking'] = params['reranking']
+
+    if 'sustainability' in params:
+        context_params['sustainability'] = params['sustainability']
+
+    logger.info("Retrieving context..")
+    try:
+        context = ir.get_context(query=query, **context_params)
+        if test:
+            retrieved_cities = ir.get_cities(context)
+        else:
+            retrieved_cities = None
+    except Exception as e:
+        exc_type, exc_obj, exc_tb = sys.exc_info()
+        logger.error(f"Error at line {exc_tb.tb_lineno} while trying to get context: {e}")
+        return None
+
+    logger.info("Retrieved context, augmenting prompt..")
+    try:
+        prompt = pg.augment_prompt(
+            query=query,
+            context=context,
+            params=context_params
+        )
+    except Exception as e:
+        exc_type, exc_obj, exc_tb = sys.exc_info()
+        logger.error(f"Error at line {exc_tb.tb_lineno} while trying to augment prompt: {e}")
+        return None
+
+    # return prompt
+
+    logger.info(f"Augmented prompt, initializing {model} and generating response..")
+    try:
+        response = tg.generate_response(model, prompt)
+    except Exception as e:
+        exc_type, exc_obj, exc_tb = sys.exc_info()
+        logger.info(f"Error at line {exc_tb.tb_lineno} while generating response: {e}")
+        return None
+
+    if test:
+        return retrieved_cities, prompt[1]['content'], response
+
+    else:
+        return response
+
+
+if __name__ == "__main__":
+    # sample_query = "I'm planning a trip in the summer and I love art, history, and visiting museums. Can you
+    # suggest " \ "some " \ "European cities? "
+    sample_query = "I'm planning a trip in July and enjoy beaches, nightlife, and vibrant cities. Recommend some " \
+                   "cities. "
+    model = "GPT-4"
+
+    pipeline_response = pipeline(
+        query=sample_query,
+        model_name=model,
+        sustainability=1
+    )
+
+    print(pipeline_response)

src/sustainability/s_fairness.py

@@ -0,0 +1,114 @@
+import sys
+import os
+import pandas as pd
+import numpy as np
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(encoding='utf-8', level=logging.DEBUG)
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.dirname(SCRIPT_DIR))
+
+from data_directories import *
+
+
+def get_popularity(destination):
+    """
+    
+    Returns the popularity score for a particular destination.
+
+    Args:
+        - destination: str
+    
+    """
+
+    parent_path = os.path.abspath(os.path.join(os.getcwd(), os.pardir))
+
+    if "src" in os.getcwd() and os.path.exists(os.path.join(parent_path, "european-city-data")):
+        popularity_path = popularity_dir.replace("../../", "../")
+    else:
+        popularity_path = popularity_dir
+
+    popularity_df = pd.read_csv(popularity_path + "popularity_scores.csv")
+
+    if not len(popularity_df[popularity_df['city'] == destination]):
+        print(f"{destination} does not have popularity data")
+        return None
+
+    return popularity_df[popularity_df['city'] == destination]['weighted_pop_score'].item()
+
+
+def get_seasonality(destination, month=None):
+    """
+    
+    Returns the seasonality score for a particular destination for a particular month. If no month is provided then
+    the best month, i.e. month of lowest seasonality is returned.
+
+    Args:
+        - destination: str
+        - month: str (default: None)
+    
+    """
+    parent_path = os.path.abspath(os.path.join(os.getcwd(), os.pardir))
+
+    if "src" in os.getcwd() and os.path.exists(os.path.join(parent_path, "european-city-data")):
+        seasonality_path = seasonality_dir.replace("../../", "../")
+    else:
+        seasonality_path = seasonality_dir
+    seasonality_df = pd.read_csv(seasonality_path + "seasonality_scores.csv")
+
+    # Check if city is present in dataframe 
+    if not len(seasonality_df[seasonality_df['city'] == destination]):
+        logger.info(f"{destination} does not have seasonality data for {month}")
+        return None, None
+
+    if month:
+        m = month.capitalize()[:3]
+    else:
+        seasonality_df['lowest_col'] = seasonality_df.loc[:, seasonality_df.columns != 'city'].idxmin(axis="columns")
+        m = seasonality_df[seasonality_df['city'] == destination]['lowest_col'].item()
+
+    # print(destination, m, seasonality_df[seasonality_df['city'] == destination][m])
+
+    return m, seasonality_df[seasonality_df['city'] == destination][m].item()
+
+
+def compute_sfairness_score(destination, month=None):
+    """
+    
+    Returns the s-fairness score for a particular destination city and (optional) month. If the destination doesn't
+    have popularity or seasonality scores, then the function returns None.
+
+    Args:
+        - destination: str
+        - month: str (default: None)
+    
+    """
+    seasonality = get_seasonality(destination, month)
+    month = seasonality[0]
+    popularity = get_popularity(destination)
+    emissions = 0
+
+    # RECHECK
+    if seasonality[1] is not None and popularity is not None:
+        s_fairness = round(0.281 * emissions + 0.334 * popularity + 0.385 * seasonality[1], 3)
+        return {
+            'month': month,
+            's-fairness': s_fairness
+        }
+    # elif popularity is not None: # => seasonality is None
+    #     s_fairness = 0.281 * emissions + 0.334 * popularity
+    # elif seasonality[1] is not None: # => popularity is None
+    #     s_fairness = 0.281 * emissions + 0.385 * seasonality[1]
+    # else: # => both are non
+    #     s_fairness =  0.281 * emissions
+    else:
+        return {
+            'month': None,
+            's-fairness': None
+        }
+
+
+if __name__ == "__main__":
+    print(compute_sfairness_score("Paris", "Oct"))

src/text_generation/model_init.py

@@ -0,0 +1,123 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from vertexai.generative_models import GenerativeModel
+
+from dotenv import load_dotenv
+from anthropic import AnthropicVertex
+import os
+from openai import OpenAI
+from src.text_generation.vertexai_setup import initialize_vertexai_params
+
+load_dotenv()
+if "OPENAI_API_KEY" in os.environ:
+    OAI_API_KEY = os.environ["OPENAI_API_KEY"]
+if "VERTEXAI_PROJECTID" in os.environ:
+    VERTEXAI_PROJECT = os.environ["VERTEXAI_PROJECTID"]
+
+
+class LLMBaseClass:
+    """
+    Base Class for text generation - user needs to provide the HF model ID while instantiating the class after which
+    the generate method can be called to generate responses
+    
+    """
+
+    def __init__(self, model_id) -> None:
+
+        match (model_id[0].lower()):
+            case "gpt-4o-mini":  # for open AI models
+                self.api_key = OAI_API_KEY
+                self.model = OpenAI(api_key=self.api_key)
+            case "claude-3-5-sonnet@20240620":  # for Claude through vertexAI
+                self.api_key = None
+                self.model = AnthropicVertex(region="europe-west1", project_id=VERTEXAI_PROJECT)
+            case "gemini-1.0-pro":
+                self.api_key = None
+                self.model = GenerativeModel(model_id[0].lower())
+            case _:  # for HF models
+                self.api_key = None
+                self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+
+                self.tokenizer.chat_template = "{%- for message in messages %}{%- if message['role'] == 'user' %}{{- " \
+                                               "bos_token + '[INST] ' + message['content'].strip() + ' [/INST]' }}{%- " \
+                                               "elif " \
+                                               "message['role'] == 'system' %}{{- '<<SYS>>\\n' + message[" \
+                                               "'content'].strip() + " \
+                                               "'\\n<</SYS>>\\n\\n' }}{%- elif message['role'] == 'assistant' %}{{- '[" \
+                                               "ASST] ' " \
+                                               "+ message['content'] + ' [/ASST]' + eos_token }}{%- endif %}{%- " \
+                                               "endfor %} "
+                # Initialize quantization to use less GPU
+                if torch.cuda.is_available():
+                    bnb_config = BitsAndBytesConfig(
+                        load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4",
+                        bnb_4bit_compute_dtype=torch.bfloat16
+                    )
+                else:
+                    bnb_config = None
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.bfloat16,
+                    device_map="auto",
+                    quantization_config=bnb_config,
+                )
+
+                self.model.generation_config.pad_token_id = self.tokenizer.pad_token_id
+
+                self.terminators = [
+                    self.tokenizer.eos_token_id,
+                    self.tokenizer.convert_tokens_to_ids("<|eot_id|>")
+                ]
+
+    def generate(self, messages):
+        match (self.model_id[0].lower()):
+            case "gpt-4o-mini":
+                completion = self.model.chat.completions.create(
+                    model=self.model_id[0],
+                    messages=messages,
+                    temperature=0.6,
+                    top_p=0.9,
+                )
+                # Return the generated content from the API response
+                return completion.choices[0].message.content
+            case "claude-3-5-sonnet@20240620" | "gemini-1.0-pro":
+                initialize_vertexai_params()
+                if "claude" in self.model_id[0].lower():
+                    message = self.model.messages.create(
+                        max_tokens=1024,
+                        model=self.model_id[0],
+                        messages=[
+                            {
+                              "role": "user",
+                              "content": messages[0]["content"],
+                            }
+                          ],
+                    )
+                    return message.content[0].text
+                else:
+                    response = self.model.generate_content(messages[0]["content"])
+                    return response
+            case _:
+                input_ids = self.tokenizer.apply_chat_template(
+                    conversation=messages,
+                    add_generation_prompt=True,
+                    return_tensors="pt",
+                    padding=True
+                ).to(self.model.device)
+
+                outputs = self.model.generate(
+                    input_ids,
+                    max_new_tokens=1024,
+                    # eos_token_id=self.terminators,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    do_sample=True,
+                    temperature=0.6,
+                    top_p=0.9,
+                )
+                response = outputs[0][input_ids.shape[-1]:]
+
+                return self.tokenizer.decode(response, skip_special_tokens=True)
+
+
+# database/wikivoyage/wikivoyage_listings.lance/data/e2940f51-d754-4b54-a688-004bdb8e7aa2.lance

src/text_generation/models.py

@@ -0,0 +1,112 @@
+from src.text_generation.model_init import LLMBaseClass
+
+
+class Llama3(LLMBaseClass):
+    """
+    Initializes a Llama3 model object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "meta-llama/Meta-Llama-3-8B"
+
+        super().__init__(self.model_id)
+
+
+class Mistral(LLMBaseClass):
+    """
+    Initializes a Mistral model object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "mistralai/Mistral-7B-v0.3"
+        super().__init__(self.model_id)
+
+
+class Gemma2(LLMBaseClass):
+    """
+    Initializes a Gemma2 model object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "google/gemma-2-9b"
+        super().__init__(self.model_id)
+
+
+class Llama3Point1(LLMBaseClass):
+    """
+    Initializes a Llama 3.1 object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "meta-llama/Meta-Llama-3.1-8B"
+        super().__init__(self.model_id)
+
+
+class Llama3Instruct(LLMBaseClass):
+    """
+    Initializes a Llama 3 Instruct object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+        super().__init__(self.model_id)
+
+
+class MistralInstruct(LLMBaseClass):
+    """
+    Initializes a Mistral Instruct object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+        super().__init__(self.model_id)
+
+
+class Llama3Point1Instruct(LLMBaseClass):
+    """
+    Initializes a Llama 3.1 Instruct object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+        super().__init__(self.model_id)
+
+
+class Phi3SmallInstruct(LLMBaseClass):
+    """
+    Initializes a Phi3-Small-Instruct object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "microsoft/Phi-3-small-128k-instruct"
+        super().__init__(self.model_id)
+
+
+class GPT4(LLMBaseClass):
+    """
+    Initializes a GPT-4 Instruct object
+    """
+
+    def __init__(self) -> None:
+        self.model_id = "gpt-4o-mini",
+        super().__init__(self.model_id)
+
+
+class Claude3Point5Sonnet(LLMBaseClass):
+    """
+       Initializes a Claude3.5 Instruct object
+       """
+
+    def __init__(self) -> None:
+        self.model_id = "claude-3-5-sonnet@20240620",
+        super().__init__(self.model_id)
+
+
+class Gemini(LLMBaseClass):
+    """
+       Initializes a Gemini Instruct object
+       """
+
+    def __init__(self) -> None:
+        self.model_id = "gemini-1.0-pro",
+        super().__init__(self.model_id)

src/text_generation/text_generation.py

@@ -0,0 +1,85 @@
+from augmentation import prompt_generation as pg
+from information_retrieval import info_retrieval as ir
+from src.text_generation.models import (
+    Llama3,
+    Mistral,
+    Gemma2,
+    Llama3Point1,
+    GPT4,
+    Claude3Point5Sonnet,
+)
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(encoding='utf-8', level=logging.DEBUG)
+
+
+def generate_response(model, prompt):
+    """
+    
+    Function that initializes the LLM class and calls the generate function.
+
+    Args:
+        - messages: list; contains the system and user prompt 
+        - model: class; the class of the llm to be initialized
+    
+    """
+
+    logger.info(f"Initializing LLM configuration for {model}")
+    llm = model()
+
+    logger.info("Generating response")
+    try:
+        response = llm.generate(prompt)
+    except Exception as e:
+        logger.error(f"Error while generating response for {model}: {e}")
+        response = 'ERROR'
+
+    return response
+
+
+def test(model):
+    context_params = {
+        'limit': 3,
+        'reranking': 0
+    }
+    # model = Llama3Point1
+
+    query = "Suggest some places to visit during winter. I like hiking, nature and the mountains and I enjoy skiing " \
+            "in winter. "
+
+    # without sustainability
+    logger.info("Retrieving context..")
+    try:
+        context = ir.get_context(query=query, **context_params)
+    except Exception as e:
+        logger.error(f"Error while trying to get context: {e}")
+        return None
+
+    logger.info("Retrieved context, augmenting prompt (without sustainability)..")
+    try:
+        without_sfairness = pg.augment_prompt(
+            query=query,
+            context=context,
+            sustainability=0,
+            params=context_params
+        )
+    except Exception as e:
+        logger.error(f"Error while trying to augment prompt: {e}")
+        return None
+
+    # return without_sfairness
+
+    logger.info(f"Augmented prompt, initializing {model} and generating response..")
+    try:
+        response = generate_response(model, without_sfairness)
+    except Exception as e:
+        logger.info(f"Error while generating response: {e}")
+        return None
+
+    return response
+
+
+if __name__ == "__main__":
+    response = test(Claude3Point5Sonnet)
+    print(response)

src/text_generation/vertexai_setup.py

@@ -0,0 +1,37 @@
+from typing import Optional
+from dotenv import load_dotenv
+from google.oauth2 import service_account
+import vertexai
+import os
+import json
+import base64
+
+load_dotenv()
+if "VERTEXAI_PROJECTID" in os.environ:
+    VERTEXAI_PROJECT = os.environ["VERTEXAI_PROJECTID"]
+
+
+def decode_service_key():
+    encoded_key = os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
+    original_service_key = json.loads(base64.b64decode(encoded_key).decode('utf-8'))
+    if original_service_key:
+        return original_service_key
+    return None
+
+
+def initialize_vertexai_params(location: Optional[str] = "us-central1"):
+
+    creds_file_name = os.getcwd() + "/.config/application_default_credentials.json"
+    print(creds_file_name)
+    if not os.path.exists(os.path.dirname(creds_file_name)):
+        credentials = decode_service_key()
+        with open(creds_file_name, 'w', encoding='utf-8') as file:
+            json.dump(credentials, file, ensure_ascii=False, indent=4)
+
+    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_file_name
+
+    service_account.Credentials.from_service_account_file(
+        filename=os.environ["GOOGLE_APPLICATION_CREDENTIALS"],
+        scopes=["https://www.googleapis.com/auth/cloud-platform"],
+    )
+    vertexai.init(project=VERTEXAI_PROJECT, location=location)

src/vectordb/create_db.py

@@ -0,0 +1,24 @@
+from src.vectordb.vectordb import *
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(encoding='utf-8', level=logging.DEBUG)
+
+
+def run():
+    logging.info("Creating database for Wikivoyage Documents")
+    try:
+        create_wikivoyage_docs_db_and_add_data()
+    except Exception as e:
+        logger.error(f"Error for Wikivoyage Documents: {e}")
+
+    logging.info("Creating database for Wikivoyage Listings")
+    try:
+        create_wikivoyage_listings_db_and_add_data()
+        print("Completed")
+    except Exception as e:
+        logger.error(f"Error for Wikivoyage Listings: {e}")
+
+
+if __name__ == "__main__":
+    run()

src/vectordb/helpers.py

@@ -0,0 +1,150 @@
+import pandas as pd
+import os
+import re
+from sentence_transformers import SentenceTransformer
+import sys
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.dirname(SCRIPT_DIR))
+
+from data_directories import *
+
+
+def create_chunks(city, country, text):
+    """
+    
+    Helper function that creates chunks given paragraph(s) of text based on implicit sections in the text. 
+
+    Args: 
+        - city: str
+        - country: str
+        - text: str; document that needs to be chunked
+
+    """
+
+    for i, line in enumerate(text):
+        if line[0] == "\n":
+            del text[i]
+
+    index = 0
+    chunks = []
+    pattern = re.compile("==")
+    ignore = re.compile("===")
+    section = 'Introduction'
+    for i, line in enumerate(text):
+        if pattern.search(line) and not ignore.search(line):
+            chunk = ''.join(text[index:i])
+            chunks.append({
+                'city': city,
+                'country': country,
+                'section': section,
+                'text': chunk,
+                # 'vector': f'city: {city}, country: {country}, section: {section}, text: {chunk}' 
+            })
+            index = i + 1
+            section = re.sub(pattern, '', line).strip()
+
+    df = pd.DataFrame(chunks)
+    return df
+
+
+def read_docs():
+    """
+    
+    Helper function that reads all of the Wikivoyage documents containing information about the city. 
+
+    """
+
+    df = pd.DataFrame()
+    cities = pd.read_csv(cities_csv)
+    for file_name in os.listdir(wikivoyage_docs_dir + "cleaned/"):
+        city = file_name.split(".")[0]
+        # print(city)
+        country = cities[cities['city'] == city]['country'].item()
+        with open(wikivoyage_docs_dir + "cleaned/" + file_name) as file:
+            text = file.readlines()
+            chunk_df = create_chunks(city, country, text)
+            df = pd.concat([df, chunk_df])
+
+    return df
+
+
+def read_listings():
+    """
+    
+    Helper function that reads the Wikivoyage listings csv containing tabular information about 144 cities. 
+
+    """
+    df = pd.read_csv(wikivoyage_listings_dir + "wikivoyage-listings-cleaned.csv")
+    cities = pd.read_csv(cities_csv)
+
+    def find_country(city):
+        return cities[cities['city'] == city]['country'].values[0]
+
+    df['country'] = df['city'].apply(find_country)
+
+    return df
+
+
+def preprocess_df(df):
+    """
+    
+    Helper function that preprocesses the dataframe containing chunks of text and removes hyperlinks and strips the \n from the text. 
+
+    Args:
+        - df: dataframe
+
+    """
+    section_counts = df['section'].value_counts()
+    sections_to_keep = section_counts[section_counts > 150].index
+    filtered_df = df[df['section'].isin(sections_to_keep)]
+
+    def preprocess_text(s):
+        s = re.sub(r'http\S+', '', s)
+        s = re.sub(r'=+', '', s)
+        s = s.strip()
+        return s
+
+    filtered_df['text'] = filtered_df['text'].apply(preprocess_text)
+
+    return filtered_df
+
+
+def compute_wv_docs_embeddings(df):
+    """
+    
+    Helper function that computes embeddings for the text. The all-MiniLM-L6-v2 embedding model is used.  
+
+    Args:
+        - df: dataframe
+
+    """
+    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    vector_dimension = model.get_sentence_embedding_dimension()
+
+    print("Computing embeddings")
+    embeddings = []
+    for i, row in df.iterrows():
+        emb = model.encode(row['combined'], show_progress_bar=True).tolist()
+        embeddings.append(emb)
+
+    print("Finished computing embeddings for wikivoyage documents.")
+    df['vector'] = embeddings
+    # df.to_csv(wv_embeddings + "wikivoyage-listings-embeddings.csv")
+    # print("Finished saving file.")
+    return df
+
+
+def embed_query(query):
+    """
+    
+    Helper function that returns the embedded query. 
+
+    Args:
+        - query: str
+    
+    """
+    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    # vector_dimension = model.get_sentence_embedding_dimension()   
+    embedding = model.encode(query).tolist()
+    return embedding

src/vectordb/lancedb_init.py

@@ -0,0 +1,37 @@
+import os
+import sys
+from lancedb.embeddings import get_registry
+from lancedb.pydantic import LanceModel, Vector
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.dirname(SCRIPT_DIR))
+
+
+model = get_registry().get("sentence-transformers").create()
+
+
+class WikivoyageDocuments(LanceModel):
+    """
+    
+    Schema definition for the Wikivoyage Documents table.
+
+    """
+    city: str = model.SourceField()
+    country: str = model.SourceField()
+    section: str = model.SourceField()
+    text: str = model.SourceField()
+    vector: Vector(model.ndims()) = model.VectorField()
+
+
+class WikivoyageListings(LanceModel):
+    """
+    
+    Schema definition for the Wikivoyage Listings table.
+
+    """
+    city: str = model.SourceField()
+    type: str = model.SourceField()
+    title: str = model.SourceField()
+    description: str = model.SourceField()
+    country: str = model.SourceField()
+    vector: Vector(model.ndims()) = model.VectorField()

src/vectordb/vectordb.py

@@ -0,0 +1,186 @@
+# from src import * 
+from src.vectordb.helpers import *
+from src.vectordb.lancedb_init import *
+import logging
+import os
+import lancedb
+from lancedb.rerankers import ColbertReranker
+
+import sys
+logger = logging.getLogger(__name__)
+from typing import Optional
+
+# db = lancedb.connect("/tmp/db")
+
+def create_wikivoyage_docs_db_and_add_data():
+    """
+    
+    Creates wikivoyage documents table and ingests data
+
+    """
+    uri = database_dir
+    current_dir = os.path.split(os.getcwd())[1]
+
+    if "src" or "tests" in current_dir: # hacky way to get the correct path
+        uri = uri.replace("../../", "../")
+
+    db = lancedb.connect(uri)
+    logger.info("Connected to DB. Reading data now...")
+    df = read_docs()
+    filtered_df = preprocess_df(df)
+    logger.info("Finished reading data, attempting to create table and ingest the data...")
+
+    db.drop_table("wikivoyage_documents", ignore_missing=True)
+    table = db.create_table("wikivoyage_documents", schema=WikivoyageDocuments)
+
+    table.add(filtered_df.to_dict('records'))
+    logger.info("Completed ingestion.")
+
+
+def create_wikivoyage_listings_db_and_add_data():
+    """
+    
+    Creates wikivoyage listings table and ingests data
+
+    """
+    uri = database_dir
+    current_dir = os.path.split(os.getcwd())[1]
+
+    if "src" or "tests" in current_dir: # hacky way to get the correct path
+        uri = uri.replace("../../", "../")
+        
+    db = lancedb.connect(uri)
+    logger.info("Connected to DB. Reading data now...")
+    df = read_listings()
+    logger.info("Finished reading data, attempting to create table and ingest the data...")
+    # filtered_df = preprocess_df(df)
+
+    db.drop_table("wikivoyage_listings", ignore_missing=True)
+    table = db.create_table("wikivoyage_listings", schema=WikivoyageListings)
+
+    table.add(df.astype('str').to_dict('records'))
+    logger.info("Completed ingestion.")
+
+
+def search_wikivoyage_docs(query: str, limit: int = 10, reranking: int = 0, run_local: Optional[bool] = False):
+    """
+    
+    Function to search the wikivoyage database an return most relevant chunked docs. 
+
+    Args: 
+        - query: str
+        - limit: number of results (default is 10)
+        - reranking: bool (0 or 1), if activated, CrossEncoderReranker is used.
+
+    """
+    if run_local:
+        uri = database_dir
+        current_dir = os.path.split(os.getcwd())[1]
+
+        if "src" or "tests" in current_dir: # hacky way to get the correct path
+            uri = uri.replace("../../", "../")
+    else:
+        uri = os.environ["BUCKET_NAME"]
+    # print(uri)
+    try:
+        db = lancedb.connect(uri)
+    except Exception as e:
+        logger.error(f"Error while connecting to DB: {e}")
+
+    logger.info("Connected to Wikivoyage DB.")
+
+    # query_embedding = embed_query(query)
+    table = db.open_table("wikivoyage_documents")
+
+    if reranking:
+        try:
+            reranker = ColbertReranker(column='text')
+            results = table.search(query) \
+                .metric('cosine') \
+                .rerank(reranker=reranker) \
+                .limit(limit) \
+                .to_list()
+        except Exception as e:
+            exc_type, exc_obj, exc_tb = sys.exc_info()
+            fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+            logger.error(f"Error while getting context: {e}, {(exc_type, fname, exc_tb.tb_lineno)}")
+
+    else:
+        try:
+            results = table.search(query) \
+                .limit(limit) \
+                .metric('cosine') \
+                .to_list()
+        except Exception as e:
+            exc_type, exc_obj, exc_tb = sys.exc_info()
+            fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+            logger.error(f"Error while getting context: {e}, {(exc_type, fname, exc_tb.tb_lineno)}")
+
+    logger.info("Found the most relevant documents.")
+    city_lists = [{"city": r['city'], "country": r['country'], "section": r['section'], "text": r['text']} for r in
+                  results]
+
+    # context = [f"city: {r['city']}, country: {r['country']}, name: {r['title']}, description: {r['description']}"
+    # for r in results]
+
+    return city_lists
+
+
+def search_wikivoyage_listings(query, cities, limit=10, reranking=0):
+    """
+    
+    Function to search the wikivoyage database an return most relevant listings, post-filtered by the recommended
+    cities provided by wikivoyage_documents table.
+
+    Args: 
+        - query: str
+        - cities: list
+        - limit: number of results (default is 10)
+        - reranking: bool (0 or 1), if activated, CrossEncoderReranker is used.
+
+    """
+    uri = database_dir
+    current_dir = os.path.split(os.getcwd())[1]
+
+    if "src" or "tests" in current_dir: # hacky way to get the correct path
+        uri = uri.replace("../../", "../")
+
+    db = lancedb.connect(uri)
+    logger.info("Connected to Wikivoyage Listings DB.")
+
+    table = db.open_table("wikivoyage_listings")
+
+    cities_filter = f"city IN {tuple(cities)}"
+
+    if reranking:
+        try:
+            reranker = ColbertReranker(column='description')
+            results = table.search(query) \
+                .where(cities_filter) \
+                .metric('cosine') \
+                .rerank(reranker=reranker) \
+                .limit(limit) \
+                .to_list()
+
+        except Exception as e:
+            exc_type, exc_obj, exc_tb = sys.exc_info()
+            fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+            logger.error(f"Error while getting context: {e}, {(exc_type, fname, exc_tb.tb_lineno)}")
+
+    else:
+        try:
+            results = table.search(query) \
+                .where(cities_filter) \
+                .metric('cosine') \
+                .limit(limit) \
+                .to_list()
+        except Exception as e:
+            exc_type, exc_obj, exc_tb = sys.exc_info()
+            fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
+            logger.error(f"Error while getting context: {e}, {(exc_type, fname, exc_tb.tb_lineno)}")
+
+    logger.info("Found the most relevant documents.")
+    city_listings = [{"city": r['city'], "country": r['country'], "type": r['type'], "title": r['title'],
+                      "description": r['description']} for r in results]
+
+    return city_listings