Spaces:
Running
Running
| from pathlib import Path | |
| import duckdb | |
| import holoviews as hv | |
| import pandas as pd | |
| import panel as pn | |
| from bokeh.models import HoverTool | |
| from bokeh.models import NumeralTickFormatter | |
| from pydantic import BaseModel, Field | |
| from langchain.callbacks.base import BaseCallbackHandler | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.llms.openai import OpenAI | |
| from langchain.output_parsers import PydanticOutputParser | |
| from langchain.pydantic_v1 import BaseModel, Field, validator | |
| from langchain.memory import ConversationBufferMemory | |
| from langchain.chains import ConversationChain | |
| from langchain.prompts import PromptTemplate | |
| pn.extension(sizing_mode="stretch_width", notifications=True) | |
| hv.extension("bokeh") | |
| INSTRUCTIONS = """ | |
| #### Name Chronicles lets you explore the history of names in the United States. | |
| - Enter a name to add to plot! | |
| - Hover over a line for stats or click for the gender distribution. | |
| - Chat with AI for inspiration or get a random name based on input criteria. | |
| - Have ideas? [Open an issue](https://github.com/ahuang11/name-chronicles/issues). | |
| """ | |
| RANDOM_NAME_QUERY = """ | |
| SELECT name, count, | |
| CASE | |
| WHEN female_percent >= 0.2 AND female_percent <= 0.8 AND male_percent >= 0.2 AND male_percent <= 0.8 THEN 'unisex' | |
| WHEN female_percent > 0.5 THEN 'female' | |
| WHEN male_percent > 0.5 THEN 'male' | |
| END AS gender | |
| FROM ( | |
| SELECT | |
| name, | |
| MAX(male + female) AS count, | |
| (SUM(female) / CAST(SUM(male + female) AS REAL)) AS female_percent, | |
| (SUM(male) / CAST(SUM(male + female) AS REAL)) AS male_percent | |
| FROM names | |
| WHERE name LIKE ? | |
| GROUP BY name | |
| ) | |
| WHERE count >= ? AND count <= ? | |
| AND gender = ? | |
| ORDER BY RANDOM() | |
| LIMIT 100 | |
| """ | |
| TOP_NAMES_WILDCARD_QUERY = """ | |
| SELECT name, SUM(male + female) as count | |
| FROM names | |
| WHERE lower(name) LIKE ? | |
| GROUP BY name | |
| ORDER BY count DESC | |
| LIMIT 10 | |
| """ | |
| TOP_NAMES_SELECT_QUERY = """ | |
| SELECT name, SUM(male + female) as count | |
| FROM names | |
| WHERE lower(name) = ? | |
| GROUP BY name | |
| ORDER BY count DESC | |
| """ | |
| DATA_QUERY = """ | |
| SELECT name, year, male, female, SUM(male + female) AS count | |
| FROM names | |
| WHERE name in ({placeholders}) | |
| GROUP BY name, year, male, female | |
| ORDER BY name, year | |
| """ | |
| MAX_LLM_COUNT = 2000 | |
| class FirstNames(BaseModel): | |
| names: list[str] = Field(description="List of first names") | |
| class StreamHandler(BaseCallbackHandler): | |
| def __init__(self, container, initial_text="", target_attr="value"): | |
| self.container = container | |
| self.text = initial_text | |
| self.target_attr = target_attr | |
| def on_llm_new_token(self, token: str, **kwargs) -> None: | |
| self.text += token | |
| setattr(self.container, self.target_attr, self.text) | |
| class NameChronicles: | |
| def __init__(self): | |
| super().__init__() | |
| self.llm_use_counter = 0 | |
| self.db_path = Path("data/names.db") | |
| # Main | |
| self.scatter_cycle = hv.Cycle("Category10") | |
| self.curve_cycle = hv.Cycle("Category10") | |
| self.label_cycle = hv.Cycle("Category10") | |
| self.holoviews_pane = pn.pane.HoloViews( | |
| min_height=675, sizing_mode="stretch_both" | |
| ) | |
| self.selection = hv.streams.Selection1D() | |
| # Sidebar | |
| # Name Widgets | |
| self.names_input = pn.widgets.TextInput(name="Name Input", placeholder="Andrew") | |
| self.names_input.param.watch(self._add_name, "value") | |
| self.names_choice = pn.widgets.MultiChoice( | |
| name="Selected Names", | |
| options=["Andrew"], | |
| solid=False, | |
| ) | |
| self.names_choice.param.watch(self._update_plot, "value") | |
| # Reset Widgets | |
| self.clear_button = pn.widgets.Button( | |
| name="Clear Names", button_style="outline", button_type="primary" | |
| ) | |
| self.clear_button.on_click( | |
| lambda event: setattr(self.names_choice, "value", []) | |
| ) | |
| self.refresh_button = pn.widgets.Button( | |
| name="Refresh Plot", button_style="outline", button_type="primary" | |
| ) | |
| self.refresh_button.on_click(self._refresh_plot) | |
| # Randomize Widgets | |
| self.name_pattern = pn.widgets.TextInput( | |
| name="Name Pattern", placeholder="*na*" | |
| ) | |
| self.count_range = pn.widgets.IntRangeSlider( | |
| name="Peak Count Range", | |
| value=(0, 100000), | |
| start=0, | |
| end=100000, | |
| step=1000, | |
| margin=(5, 20), | |
| ) | |
| self.gender_select = pn.widgets.RadioButtonGroup( | |
| name="Gender", | |
| options=["Female", "Unisex", "Male"], | |
| button_style="outline", | |
| button_type="primary", | |
| ) | |
| randomize_name = pn.widgets.Button( | |
| name="Get Name", button_style="outline", button_type="primary" | |
| ) | |
| randomize_name.param.watch(self._randomize_name, "clicks") | |
| self.randomize_pane = pn.Card( | |
| self.name_pattern, | |
| self.count_range, | |
| self.gender_select, | |
| randomize_name, | |
| title="Get Random Name", | |
| collapsed=True, | |
| ) | |
| # AI Widgets | |
| self.chat_interface = pn.chat.ChatInterface( | |
| show_button_name=False, | |
| callback=self._prompt_ai, | |
| height=500, | |
| styles={"background": "white"}, | |
| disabled=True, | |
| ) | |
| self.chat_interface.send( | |
| value=( | |
| "Ask me about name suggestions or their history! " | |
| "To add suggested names, click the button below!" | |
| ), | |
| user="System", | |
| respond=False, | |
| ) | |
| self.parse_ai_button = pn.widgets.Button( | |
| name="Parse and Add Names", | |
| button_style="outline", | |
| button_type="primary", | |
| disabled=True, | |
| ) | |
| self.last_ai_output = None | |
| pn.state.onload(self._initialize_database) | |
| # Database Methods | |
| def _initialize_database(self): | |
| """ | |
| Initialize database with data from the Social Security Administration. | |
| """ | |
| self.conn = duckdb.connect(":memory:") | |
| df = pd.concat( | |
| [ | |
| pd.read_csv( | |
| path, | |
| header=None, | |
| names=["state", "gender", "year", "name", "count"], | |
| ) | |
| for path in Path("data").glob("*.TXT") | |
| ] | |
| ) | |
| df_processed = ( | |
| df.groupby(["gender", "year", "name"], as_index=False)[["count"]] | |
| .sum() | |
| .pivot(index=["name", "year"], columns="gender", values="count") | |
| .reset_index() | |
| .rename(columns={"F": "female", "M": "male"}) | |
| .fillna(0) | |
| ) | |
| self.conn.execute("DROP TABLE IF EXISTS names") | |
| self.conn.execute("CREATE TABLE names AS SELECT * FROM df_processed") | |
| if self.names_choice.value == []: | |
| self.names_choice.value = ["Andrew"] | |
| else: | |
| self.names_choice.param.trigger("value") | |
| self.main.objects = [self.holoviews_pane] | |
| # Start AI | |
| self.callback_handler = pn.chat.langchain.PanelCallbackHandler( | |
| self.chat_interface | |
| ) | |
| self.chat_openai = ChatOpenAI( | |
| max_tokens=75, | |
| streaming=True, | |
| callbacks=[self.callback_handler], | |
| ) | |
| self.openai = OpenAI(max_tokens=75) | |
| memory = ConversationBufferMemory() | |
| self.conversation_chain = ConversationChain( | |
| llm=self.chat_openai, memory=memory, callbacks=[self.callback_handler] | |
| ) | |
| self.chat_interface.disabled = False | |
| self.parse_ai_button.on_click(self._parse_ai_output) | |
| self.pydantic_parser = PydanticOutputParser(pydantic_object=FirstNames) | |
| self.prompt_template = PromptTemplate( | |
| template="{format_instructions}\n{input}\n", | |
| input_variables=["input"], | |
| partial_variables={"format_instructions": self.pydantic_parser.get_format_instructions()}, | |
| ) | |
| def _query_names(self, names): | |
| """ | |
| Query the database for the given name. | |
| """ | |
| dfs = [] | |
| for name in names: | |
| if "*" in name or "%" in name: | |
| name = name.replace("*", "%") | |
| top_names_query = TOP_NAMES_WILDCARD_QUERY | |
| else: | |
| top_names_query = TOP_NAMES_SELECT_QUERY | |
| top_names = ( | |
| self.conn.execute(top_names_query, [name.lower()]) | |
| .fetch_df()["name"] | |
| .tolist() | |
| ) | |
| if len(top_names) == 0: | |
| pn.state.notifications.info(f"No names found matching {name!r}") | |
| continue | |
| data_query = DATA_QUERY.format( | |
| placeholders=", ".join(["?"] * len(top_names)) | |
| ) | |
| df = self.conn.execute(data_query, top_names).fetch_df() | |
| dfs.append(df) | |
| if len(dfs) > 0: | |
| self.df = pd.concat(dfs).drop_duplicates( | |
| subset=["name", "year", "male", "female"] | |
| ) | |
| else: | |
| self.df = pd.DataFrame(columns=["name", "year", "male", "female"]) | |
| # Widget Methods | |
| def _randomize_name(self, event): | |
| name_pattern = self.name_pattern.value.lower() | |
| if not name_pattern: | |
| name_pattern = "%" | |
| else: | |
| name_pattern = name_pattern.replace("*", "%") | |
| if not name_pattern.startswith("%"): | |
| name_pattern = name_pattern.title() | |
| count_range = self.count_range.value | |
| gender_select = self.gender_select.value.lower() | |
| random_names = ( | |
| self.conn.execute( | |
| RANDOM_NAME_QUERY, [name_pattern, *count_range, gender_select] | |
| ).fetch_df()["name"] | |
| .tolist() | |
| ) | |
| print(len(random_names)) | |
| if random_names: | |
| for i in range(len(random_names)): | |
| random_name = random_names[i] | |
| if random_name in self.names_choice.value: | |
| continue | |
| self.names_input.value = random_name | |
| break | |
| else: | |
| pn.state.notifications.info( | |
| "All names matching the criteria are already added!" | |
| ) | |
| else: | |
| pn.state.notifications.info("No names found matching the criteria!") | |
| def _add_only_unique_names(self, names): | |
| value = self.names_choice.value.copy() | |
| options = self.names_choice.options.copy() | |
| for name in names: | |
| if " " in name: | |
| name = name.split(" ", 1)[0] | |
| if name not in options: | |
| options.append(name) | |
| if name not in value: | |
| value.append(name) | |
| self.names_choice.param.update( | |
| options=options, | |
| value=value, | |
| ) | |
| def _add_name(self, event): | |
| name = event.new.strip().title() | |
| self.names_input.value = "" | |
| if not name: | |
| return | |
| elif name in self.names_choice.options and name in self.names_choice.value: | |
| pn.state.notifications.info(f"{name!r} already added!") | |
| return | |
| elif len(self.names_choice.value) > 10: | |
| pn.state.notifications.info( | |
| "Maximum of 10 names allowed; please remove some first!" | |
| ) | |
| return | |
| self._add_only_unique_names([name]) | |
| async def _prompt_ai(self, contents, user, instance): | |
| if self.llm_use_counter >= MAX_LLM_COUNT: | |
| pn.state.notifications.info( | |
| "Sorry, all the available AI credits have been used!" | |
| ) | |
| return | |
| prompt = ( | |
| f"One sentence reply to {contents!r} or concisely suggest other relevant names; " | |
| f"if no name is provided use {self.names_choice.value[-1]!r}." | |
| ) | |
| print(prompt) | |
| self.last_ai_output = await self.conversation_chain.apredict( | |
| input=prompt, | |
| callbacks=[self.callback_handler], | |
| ) | |
| self.parse_ai_button.disabled = False | |
| self.llm_use_counter += 1 | |
| async def _parse_ai_output(self, _): | |
| if self.llm_use_counter >= MAX_LLM_COUNT: | |
| pn.state.notifications.info( | |
| "Sorry, all the available AI credits have been used!" | |
| ) | |
| return | |
| if self.last_ai_output is None: | |
| pn.state.notifications.info("No available AI output to parse!") | |
| return | |
| try: | |
| names_prompt = self.prompt_template.format_prompt(input=self.last_ai_output).to_string() | |
| names_text = await self.openai.apredict(names_prompt) | |
| new_names = (await self.pydantic_parser.aparse(names_text)).names | |
| print(new_names) | |
| self._add_only_unique_names(new_names) | |
| except Exception: | |
| pn.state.notifications.error("Failed to parse AI output.") | |
| finally: | |
| self.last_ai_output = None | |
| # Plot Methods | |
| def _click_plot(self, index): | |
| gender_nd_overlay = hv.NdOverlay(kdims=["Gender"]) | |
| if not index: | |
| return hv.NdOverlay( | |
| { | |
| "curve": self._curve_nd_overlay, | |
| "scatter": self._scatter_nd_overlay, | |
| "label": self._label_nd_overlay, | |
| } | |
| ) | |
| name = self._name_indices[index[0]] | |
| df_name = self.df.loc[self.df["name"] == name].copy() | |
| df_name["female"] += df_name["male"] | |
| gender_nd_overlay["Male"] = hv.Area( | |
| df_name, ["year"], ["male"], label="Male" | |
| ).opts(alpha=0.3, color="#add8e6", line_alpha=0) | |
| gender_nd_overlay["Female"] = hv.Area( | |
| df_name, ["year"], ["male", "female"], label="Female" | |
| ).opts(alpha=0.3, color="#ffb6c1", line_alpha=0) | |
| return hv.NdOverlay( | |
| { | |
| "curve": self._curve_nd_overlay[[index[0]]], | |
| "scatter": self._scatter_nd_overlay, | |
| "label": self._label_nd_overlay[[index[0]]].opts(text_color="black"), | |
| "gender": gender_nd_overlay, | |
| }, | |
| kdims=["Gender"], | |
| ).opts(legend_position="top_left") | |
| def _update_plot(self, event): | |
| names = event.new | |
| print(names) | |
| self._query_names(names) | |
| self._scatter_nd_overlay = hv.NdOverlay() | |
| self._curve_nd_overlay = hv.NdOverlay(kdims=["Name"]).opts( | |
| gridstyle={"xgrid_line_width": 0}, | |
| show_grid=True, | |
| fontscale=1.28, | |
| xlabel="Year", | |
| ylabel="Count", | |
| yformatter=NumeralTickFormatter(format="0.0a"), | |
| legend_limit=0, | |
| padding=(0.2, 0.05), | |
| title="Name Chronicles", | |
| responsive=True, | |
| ) | |
| self._label_nd_overlay = hv.NdOverlay(kdims=["Name"]) | |
| hover_tool = HoverTool( | |
| tooltips=[("Name", "@name"), ("Year", "@year"), ("Count", "@count")], | |
| ) | |
| self._name_indices = {} | |
| for i, (name, df_name) in enumerate(self.df.groupby("name")): | |
| df_name_total = df_name.groupby( | |
| ["name", "year", "male", "female"], as_index=False | |
| )["count"].sum() | |
| df_name_total["male"] = df_name_total["male"] / df_name_total["count"] | |
| df_name_total["female"] = df_name_total["female"] / df_name_total["count"] | |
| df_name_peak = df_name.loc[[df_name["count"].idxmax()]] | |
| df_name_peak[ | |
| "label" | |
| ] = f'{df_name_peak["name"].item()} ({df_name_peak["year"].item()})' | |
| hover_tool = HoverTool( | |
| tooltips=[ | |
| ("Name", "@name"), | |
| ("Year", "@year"), | |
| ("Count", "@count{(0a)}"), | |
| ("Male", "@male{(0%)}"), | |
| ("Female", "@female{(0%)}"), | |
| ], | |
| ) | |
| self._scatter_nd_overlay[i] = hv.Scatter( | |
| df_name_total, ["year"], ["count", "male", "female", "name"], label=name | |
| ).opts( | |
| color=self.scatter_cycle, | |
| size=4, | |
| alpha=0.15, | |
| marker="y", | |
| tools=["tap", hover_tool], | |
| line_width=3, | |
| show_legend=False, | |
| ) | |
| self._curve_nd_overlay[i] = hv.Curve( | |
| df_name_total, ["year"], ["count"], label=name | |
| ).opts( | |
| color=self.curve_cycle, | |
| tools=["tap"], | |
| line_width=3, | |
| ) | |
| self._label_nd_overlay[i] = hv.Labels( | |
| df_name_peak, ["year", "count"], ["label"], label=name | |
| ).opts( | |
| text_align="right", | |
| text_baseline="bottom", | |
| text_color=self.label_cycle, | |
| ) | |
| self._name_indices[i] = name | |
| self.selection.source = self._curve_nd_overlay | |
| if len(self._name_indices) == 1: | |
| self.selection.update(index=[0]) | |
| else: | |
| self.selection.update(index=[]) | |
| self.dynamic_map = hv.DynamicMap( | |
| self._click_plot, kdims=[], streams=[self.selection] | |
| ).opts(responsive=True) | |
| self._refresh_plot() | |
| def _refresh_plot(self, event=None): | |
| self.holoviews_pane.object = self.dynamic_map.clone() | |
| def view(self): | |
| reset_row = pn.Row(self.clear_button, self.refresh_button) | |
| data_url = pn.pane.Markdown( | |
| "<center>Data from the <a href='https://www.ssa.gov/oact/babynames/limits.html' " | |
| "target='_blank'>U.S. Social Security Administration</a></center>", | |
| align="end", | |
| ) | |
| sidebar = pn.Column( | |
| INSTRUCTIONS, | |
| self.names_input, | |
| self.names_choice, | |
| reset_row, | |
| pn.layout.Divider(), | |
| self.chat_interface, | |
| self.parse_ai_button, | |
| self.randomize_pane, | |
| data_url, | |
| ) | |
| self.main = pn.Column( | |
| pn.widgets.StaticText( | |
| value="Loading, this may take a few seconds...", | |
| sizing_mode="stretch_both", | |
| ), | |
| ) | |
| template = pn.template.FastListTemplate( | |
| sidebar_width=500, | |
| sidebar=[sidebar], | |
| main=[self.main], | |
| title="Name Chronicles", | |
| theme="dark", | |
| ) | |
| return template | |
| NameChronicles().view().servable() | |