Extraction files

app.py

@@ -0,0 +1,232 @@
+import logging
+import os
+from typing import Any
+
+import pandas as pd
+import streamlit as st
+from countryinfo import CountryInfo
+from dotenv import load_dotenv
+
+from common import HintType, configs, get_distance
+from hint import AudioHint, ImageHint, TextHint
+
+
+def setup_models(_cache: Any, configs: dict) -> None:
+    """Setups all hint models.
+
+    Args:
+        _cache (st.session_state): Streamlit cache object
+        configs (dict): Configurations used by the models
+    """
+    for model_type in _cache["hint_types"]:
+        if _cache["model"][model_type] is None:
+            if model_type == HintType.TEXT.value:
+                _cache["model"][model_type] = setup_text_hint(configs)
+            elif model_type == HintType.IMAGE.value:
+                _cache["model"][model_type] = setup_image_hint(configs)
+            elif model_type == HintType.AUDIO.value:
+                _cache["model"][model_type] = setup_audio_hint(configs)
+
+
+@st.cache_resource()
+def setup_text_hint(configs: dict) -> TextHint:
+    """Setups the text hint model.
+
+    Args:
+        configs (dict): Configurations used by the model
+
+    Returns:
+        TextHint: Hint model
+    """
+    with st.spinner("Loading text model..."):
+        model_configs = configs["local"][HintType.TEXT.value.lower()]
+        model_configs["hf_access_token"] = os.environ["HF_ACCESS_TOKEN"]
+        textHint = TextHint(configs=model_configs)
+        textHint.initialize()
+    return textHint
+
+
+@st.cache_resource()
+def setup_image_hint(configs: dict) -> ImageHint:
+    """Setups the image hint model.
+
+    Args:
+        configs (dict): Configurations used by the model
+
+    Returns:
+        ImageHint: Hint model
+    """
+    with st.spinner("Loading image model..."):
+        model_configs = configs["local"][HintType.IMAGE.value.lower()]
+        imageHint = ImageHint(configs=model_configs)
+        imageHint.initialize()
+    return imageHint
+
+
+@st.cache_resource()
+def setup_audio_hint(configs: dict) -> AudioHint:
+    """Setups the audio hint model.
+
+    Args:
+        configs (dict): Configurations used by the model
+
+    Returns:
+        AudioHint: Hint model
+    """
+    with st.spinner("Loading audio model..."):
+        model_configs = configs["local"][HintType.AUDIO.value.lower()]
+        audioHint = AudioHint(configs=model_configs)
+        audioHint.initialize()
+    return audioHint
+
+
+@st.cache_resource()
+def get_country_list() -> pd.DataFrame:
+    """Builds a database of countries and metadata.
+
+    Returns:
+        pd.DataFrame: Country database
+    """
+    country_list = list(CountryInfo().all().keys())
+
+    country_df = {}
+    for country in country_list:
+        try:
+            area = CountryInfo(country).area()
+            country_df[country] = area
+        except:
+            pass
+
+    country_df = pd.DataFrame(country_df.items(), columns=["country", "area"])
+    return country_df
+
+
+def pick_country(country_df: pd.DataFrame) -> str:
+    """Selects a country, the probability of each country is related to its area size.
+
+    Args:
+        country_df (pd.DataFrame): Database of country and their metadata
+
+    Returns:
+        str: The selected country
+    """
+    country = country_df.sample(n=1, weights="area")["country"].iloc[0]
+    return country
+
+
+def reset_cache() -> None:
+    """Reset the Streamlit APP cache."""
+    country_df = get_country_list()
+    st.session_state["country_list"] = country_df["country"].values.tolist()
+    st.session_state["country"] = pick_country(country_df)
+    st.session_state["hint_types"] = []
+    st.session_state["n_hints"] = 1
+    st.session_state["game_started"] = False
+    st.session_state["model"] = {
+        HintType.TEXT.value: None,
+        HintType.IMAGE.value: None,
+        HintType.AUDIO.value: None,
+    }
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+st.set_page_config(
+    page_title="Gen AI GeoGuesser",
+    page_icon="🌎",
+)
+
+if not st.session_state:
+    load_dotenv()
+    reset_cache()
+
+st.title("Generative AI GeoGuesser 🌎")
+
+st.markdown("### Guess the country based on hints generated by AI")
+
+col1, col2 = st.columns([2, 1])
+
+with col1:
+    st.session_state["hint_types"] = st.multiselect(
+        "Chose which hint types you want",
+        [x.value for x in HintType],
+        default=st.session_state["hint_types"],
+    )
+
+with col2:
+    st.session_state["n_hints"] = st.slider(
+        "Number of hints",
+        min_value=1,
+        max_value=5,
+        value=st.session_state["n_hints"],
+    )
+
+start_btn = st.button("Start game")
+
+if start_btn:
+    if not st.session_state["hint_types"]:
+        st.error("Pick at least one hint type")
+        reset_cache()
+    else:
+        print(f'Chosen country "{st.session_state["country"]}"')
+
+        setup_models(st.session_state, configs)
+
+        for hint_type in st.session_state["hint_types"]:
+            with st.spinner(f"Generating {hint_type} hint..."):
+                st.session_state["model"][hint_type].generate_hint(
+                    st.session_state["country"],
+                    st.session_state["n_hints"],
+                )
+
+        st.session_state["game_started"] = True
+
+if st.session_state["game_started"]:
+    game_col1, game_col2, game_col3 = st.columns([2, 1, 1])
+
+    with game_col1:
+        guess = st.selectbox("Country guess", ([""] + st.session_state["country_list"]))
+    with game_col2:
+        guess_btn = st.button("Make a guess")
+    with game_col3:
+        reset_btn = st.button("Reset game")
+
+    if guess_btn:
+        if st.session_state["country"] == guess:
+            st.success("Correct guess you won!")
+            st.balloons()
+        else:
+            if guess:
+                country_latlong = CountryInfo(st.session_state["country"]).latlng()
+                guess_latlong = CountryInfo(guess).latlng()
+                distance = int(get_distance(country_latlong, guess_latlong))
+                st.error(
+                    f"""
+                    Wrong guess, you missed the correct country by {distance} KM.
+                    The correct answer was {st.session_state["country"]}.
+                    """
+                )
+            else:
+                st.error("Pick a country.")
+
+    if reset_btn:
+        reset_cache()
+
+if st.session_state["game_started"]:
+    tabs = st.tabs([f"{x} hint" for x in st.session_state["hint_types"]])
+
+    for tab_idx, tab in enumerate(tabs):
+        hint_type = st.session_state["hint_types"][tab_idx]
+        with tab:
+            if st.session_state["model"][hint_type]:
+                for hint_idx, hint in enumerate(
+                    st.session_state["model"][hint_type].hints
+                ):
+                    st.markdown(f"#### Hint #{hint_idx+1}")
+                    if hint_type == HintType.TEXT.value:
+                        st.write(hint["text"])
+                    elif hint_type == HintType.IMAGE.value:
+                        st.image(hint["image"])
+                    elif hint_type == HintType.AUDIO.value:
+                        st.audio(hint["audio"], sample_rate=hint["sample_rate"])

common.py

@@ -0,0 +1,55 @@
+import logging
+import pprint
+from enum import Enum
+from math import acos, cos, radians, sin
+
+import yaml
+
+
+def parse_configs(configs_path: str) -> dict:
+    """Parse configs from the YAML file.
+
+    Args:
+        configs_path (str): Path to the YAML file
+
+    Returns:
+        dict: Parsed configs
+    """
+    configs = yaml.safe_load(open(configs_path, "r"))
+    logger.info(f"Configs: {pprint.pformat(configs)}")
+    return configs
+
+
+def get_distance(source_country: list[float], target_country: list[float]) -> float:
+    """Calculate the distance between two countries.
+
+    Args:
+        source_country (list[float]): Source country coordinates
+        target_country (list[float]): Target country coordinates
+
+    Returns:
+        float: Distance in KM
+    """
+    source_lat = radians(source_country[0])
+    source_long = radians(source_country[1])
+    target_lat = radians(target_country[0])
+    target_long = radians(target_country[1])
+    dist = 6371.01 * acos(
+        sin(source_lat) * sin(target_lat)
+        + cos(source_lat) * cos(target_lat) * cos(source_long - target_long)
+    )
+    return dist
+
+
+class HintType(Enum):
+    AUDIO = "Audio"
+    TEXT = "Text"
+    IMAGE = "Image"
+
+
+CONFIGS_PATH = "configs.yaml"
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__file__)
+
+configs = parse_configs(CONFIGS_PATH)

hint.py

@@ -0,0 +1,149 @@
+import abc
+import logging
+import re
+from typing import Any
+
+import torch
+from diffusers import AudioLDM2Pipeline, AutoPipelineForText2Image
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+SAMPLE_RATE = 16000
+
+
+class BaseHint(BaseModel, abc.ABC):
+    configs: dict
+    hints: list = []
+    model: Any = None
+
+    @abc.abstractmethod
+    def initialize(self):
+        """Initialize the hint model."""
+        pass
+
+    @abc.abstractmethod
+    def generate_hint(self, country: str, n_hints: int):
+        """Generate hints.
+
+        Args:
+            country (str): Country name used to base the hint
+            n_hints (int): Number of hints that will be generated
+        """
+        pass
+
+
+class TextHint(BaseHint):
+    tokenizer: Any = None
+
+    def initialize(self):
+        logger.info(
+            f"""Initializing text hint with model '{self.configs["model_id"]}'"""
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.configs["model_id"],
+            token=self.configs["hf_access_token"],
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.configs["model_id"],
+            torch_dtype=torch.float16,
+            token=self.configs["hf_access_token"],
+        ).to(self.configs["device"])
+        logger.info("Initialization finisehd")
+
+    def generate_hint(self, country: str, n_hints: int):
+        logger.info(f"Generating '{n_hints}' text hints")
+
+        generation_config = GenerationConfig(
+            do_sample=True,
+            max_new_tokens=self.configs["max_output_tokens"],
+            top_k=self.configs["top_k"],
+            top_p=self.configs["top_p"],
+            temperature=self.configs["temperature"],
+        )
+
+        prompt = [
+            f'Describe the country "{country}" without mentioning its name\n'
+            for _ in range(n_hints)
+        ]
+        input_ids = self.tokenizer(prompt, return_tensors="pt")
+        text_hints = self.model.generate(
+            **input_ids.to(self.configs["device"]),
+            generation_config=generation_config,
+        )
+
+        for idx, text_hint in enumerate(text_hints):
+            text_hint = (
+                self.tokenizer.decode(text_hint, skip_special_tokens=True)
+                .strip()
+                .replace(prompt[idx], "")
+                .strip()
+            )
+            text_hint = re.sub(
+                re.escape(country), "***", text_hint, flags=re.IGNORECASE
+            )
+
+            self.hints.append({"text": text_hint})
+
+        logger.info(f"Text hints '{n_hints}' successfully generated")
+
+
+class ImageHint(BaseHint):
+    def initialize(self):
+        logger.info(
+            f"""Initializing image hint with model '{self.configs["model_id"]}'"""
+        )
+        self.model = AutoPipelineForText2Image.from_pretrained(
+            self.configs["model_id"],
+            # torch_dtype=torch.float16,
+            variant="fp16",
+        ).to(self.configs["device"])
+        logger.info("Initialization finisehd")
+
+    def generate_hint(self, country: str, n_hints: int):
+        logger.info(f"Generating '{n_hints}' image hints")
+        prompt = [f"An image related to the country {country}" for _ in range(n_hints)]
+        img_hints = self.model(
+            prompt=prompt,
+            num_inference_steps=self.configs["num_inference_steps"],
+            guidance_scale=self.configs["guidance_scale"],
+        ).images
+        self.hints = [{"image": img_hint} for img_hint in img_hints]
+        logger.info(f"Image hints '{n_hints}' successfully generated")
+
+
+class AudioHint(BaseHint):
+    def initialize(self):
+        logger.info(
+            f"""Initializing audio hint with model '{self.configs["model_id"]}'"""
+        )
+        self.model = AudioLDM2Pipeline.from_pretrained(
+            self.configs["model_id"],
+            # torch_dtype=torch.float16,  # Not working with MacOS
+        ).to(self.configs["device"])
+        logger.info("Initialization finisehd")
+
+    def generate_hint(self, country: str, n_hints: int):
+        logger.info(f"Generating '{n_hints}' audio hints")
+        prompt = f"A sound that resembles the country of {country}"
+        negative_prompt = "Low quality"
+
+        audio_hints = self.model(
+            prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=self.configs["num_inference_steps"],
+            audio_length_in_s=self.configs["audio_length_in_s"],
+            num_waveforms_per_prompt=n_hints,
+        ).audios
+
+        for audio_hint in audio_hints:
+            self.hints.append(
+                {
+                    "audio": audio_hint,
+                    "sample_rate": SAMPLE_RATE,
+                }
+            )
+        logger.info(f"Audio hints '{n_hints}' successfully generated")