"
- html_output += f"There are {num_duplicated_rows} duplicated rows in the data. "
- remove_duplicates = duplicate_rows.get("remove_duplicates", False)
- html_output += f"The duplicated rows {'have' if remove_duplicates else 'have not'} been removed. "
- html_output += f""
-
- index_column_names = index_columns.get("index_column_names", [])
- if index_column_names:
- html_output += "
Index Columns
"
- html_output += f"There are {len(index_column_names)} columns that look like indices: {', '.join(index_column_names)} "
- remove_index_columns = index_columns.get("remove_columns", [])
- html_output += f"{len(remove_index_columns)} index columns {'have' if remove_index_columns else 'have not'} been removed. "
- html_output += f""
-
- missing_column_names = missing_columns.get("missing_column_names", [])
- if missing_column_names:
- html_output += "
Missing Columns
"
- html_output += f"There are {len(missing_column_names)} columns whose values are all missing: {', '.join(missing_column_names)} "
- remove_missing_columns = missing_columns.get("remove_columns", [])
- html_output += f"{len(remove_missing_columns)} missing columns {'have' if remove_missing_columns else 'have not'} been removed. "
- html_output += f""
-
- duplicated_column_names = duplicate_columns.get("duplicated_column_names", [])
- if duplicated_column_names:
- html_output += "
Duplicate Columns
"
- html_output += f"There are {len(duplicated_column_names)} groups of columns that have the same names: "
- html_output += ""
- for duplicate_column in duplicated_column_names:
- html_output += f"
{len(duplicate_column)} columns named {duplicate_column[0]}
"
- html_output += f"There are {len(x_y_column_names)} pairs of columns that look like the result of a merge: "
- html_output += ""
- for x_y_column in x_y_column_names:
- html_output += f"
"
- html_output += f"There are {len(invalid_data_examples)} columns contains values of invalid data type: "
- html_output += ""
- for column_name, example_data in invalid_data_examples.items():
- html_output += f"
{column_name} is of type {example_data['data_type']} "
- html_output += f" Examples of invalid data: {example_data['invalid_rows']}
"
- html_output += ""
- html_output += f"Rows with invalid data have been removed."
- html_output += f""
-
- return html_output
-
-def select_invalid_data_type(df, column, data_type):
- data_type = data_type.upper()
-
-
- def is_uuid(val):
- try:
- uuid.UUID(str(val))
- return True
- except ValueError:
- return False
-
- if data_type == 'NULL':
- mask = ~df[column].isnull()
- elif data_type in ['INTEGER', 'SMALLINT', 'TINYINT', 'BIGINT', 'HUGEINT']:
- numeric_mask = pd.to_numeric(df[column], errors='coerce').notnull()
- integer_mask = df[column].astype(str).str.isdigit()
- mask = ~(numeric_mask & integer_mask)
- elif data_type in ['DOUBLE', 'REAL', 'DECIMAL']:
- numeric_mask = pd.to_numeric(df[column], errors='coerce').notnull()
- null_mask = df[column].isnull()
- mask = ~(numeric_mask | null_mask)
- elif data_type == 'BOOLEAN':
- mask = ~df[column].isin([True, False, 1, 0, 'True', 'False', 'true', 'false'])
- elif data_type in data_types['VARCHAR']:
- mask = df[column].astype(str).isnull()
- elif data_type == 'DATE':
- mask = ~df[column].apply(lambda x: isinstance(x, (datetime.date, datetime.datetime, datetime.time)))
- elif data_type in ['TIME', 'TIMESTAMP', 'TIMESTAMP WITH TIME ZONE']:
- mask = ~df[column].apply(lambda x: isinstance(x, (datetime.date, datetime.datetime, datetime.time)))
- elif data_type == 'BLOB':
- mask = df[column].apply(lambda x: not isinstance(x, (bytes, bytearray)))
- elif data_type == 'UUID':
- mask = ~df[column].apply(is_uuid)
- else:
- raise ValueError(f"Data type {data_type} is not supported.")
-
- return mask
-
-
-
-def sql_cleaner(sql):
- sql = sql.replace('`', '"')
- return sql
-
-def clean_table_name(table_name):
- if not table_name[0].isalpha() and table_name[0] != "_":
- table_name = "_" + table_name
-
- return sanitize_table_name(table_name)
-
-def clean_column_name(column_name):
- reserved_words = {
- "SELECT", "FROM", "WHERE", "JOIN", "ON", "CREATE", "DROP", "TABLE",
- "INSERT", "UPDATE", "DELETE", "ALTER", "CAST", "EXECUTE", "ORDER", "GROUP", "BY"
- }
-
- if column_name.upper() in reserved_words:
- column_name += "_"
-
- return clean_table_name(column_name)
-
-def sanitize_table_name(table_name):
- return re.sub(r'[^a-zA-Z0-9_]', '_', table_name)
-
-def create_select_widget(options, callback):
- list_widget = widgets.Select(
- options=options,
- description='Choices:',
- disabled=False
- )
-
- button = widgets.Button(description="Submit")
-
- def on_button_clicked(b):
- callback(list_widget.value)
-
- button.on_click(on_button_clicked)
-
- display(list_widget, button)
-
-
-def identify_longitude_latitude(source_table_description):
-
- template = f"""Below are summary of the table. The attributes are highlighted in **bold**.
-Table: {source_table_description}.
-
-Task: Identify the pairs of longitude/latitude attribute names.
-Respond in JSON format:
-```json
-[
-{{
- "longitude_name": "attribute_1" (case sensitive),
- "latitude_name": "attribute_2
-}},
-...
-]
-```"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- summary = response['choices'][0]['message']['content']
- assistant_message = response['choices'][0]['message']
- messages.append(assistant_message)
-
- for message in messages:
- write_log(message['content'])
- write_log("---------------------")
-
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
- def verify_json_result(json_code):
- if not isinstance(json_code, list):
- raise ValueError("JSON code should be a list")
-
- for item in json_code:
- if not isinstance(item, dict):
- raise ValueError("Each item in the list should be a dictionary")
-
- required_keys = ["longitude_name", "latitude_name"]
- for key in required_keys:
- if key not in item:
- raise ValueError(f"Key '{key}' is missing in one of the items")
-
- if not isinstance(item[key], str):
- raise ValueError(f"Value for '{key}' should be a string")
-
- verify_json_result(json_code)
- return json_code
-
-
-def display_longitude_latitude(json_code, full_list, call_back):
- print(f"🤓 We have identified attributes for longitude/latitude. Please select one:")
-
- radio_options = [f"Longitude: \"{pair['longitude_name']}\", Latitude: \"{pair['latitude_name']}\"" for pair in json_code]
-
- if len(radio_options) == 0:
- print(f"🙁 There doesn't seem to be a pair of attributes for longitude/latitude")
- print(f"🤓 You can try to first extract these attributes")
- print(f"😊 GeoEncoding is under development. Please send a feature request! ")
-
-
- radio_options.append('Manual Selection:')
-
-
- selection_radio_buttons = widgets.RadioButtons(
- options=radio_options,
- disabled=False,
- layout=widgets.Layout(width='600px')
- )
-
- longitude_label = widgets.HTML(value='Longitude:')
- latitude_label = widgets.HTML(value='Latitude:')
-
-
- custom_longitude_dropdown = widgets.Dropdown(
- options=full_list,
- layout=widgets.Layout(width='15%')
- )
- custom_latitude_dropdown = widgets.Dropdown(
- options=full_list,
- layout=widgets.Layout(width='15%')
- )
-
- hbox_longitude = widgets.HBox([longitude_label, custom_longitude_dropdown,latitude_label, custom_latitude_dropdown])
-
- submit_button = widgets.Button(description="Submit")
-
- def on_submit_button_clicked(b):
- clear_output(wait=True)
-
- selected_index = selection_radio_buttons.index
- if selected_index < len(json_code):
- selected_pair = json_code[selected_index]
- call_back(selected_pair['longitude_name'], selected_pair['latitude_name'])
- else:
- call_back(custom_longitude_dropdown.value, custom_latitude_dropdown.value)
-
- submit_button.on_click(on_submit_button_clicked)
-
- display(selection_radio_buttons, hbox_longitude, submit_button)
-
- def on_selection_change(change):
- if change['new'] == 'Manual Selection:':
- hbox_longitude.layout.display = 'flex'
- else:
- hbox_longitude.layout.display = 'none'
-
- selection_radio_buttons.observe(on_selection_change, names='value')
- on_selection_change({'new': selection_radio_buttons.value})
-
-
-
-
-def create_selection_grid(columns1, columns2, table1_name, table2_name, call_back_func):
- table1_label = widgets.Label(value=table1_name)
- table2_label = widgets.Label(value=table2_name)
-
- table1_selectors = [widgets.Checkbox() for col in columns1]
- table2_selectors = [widgets.Checkbox() for col in columns2]
-
- def on_submit_clicked(b):
- selected_table1_indices = [i for i, selected in enumerate(table1_selectors) if selected.value]
- selected_table2_indices = [i for i, selected in enumerate(table2_selectors) if selected.value]
- call_back_func(selected_table1_indices, selected_table2_indices)
-
- submit_button = widgets.Button(description="Submit")
- submit_button.on_click(on_submit_clicked)
-
- grid = widgets.GridspecLayout(2 + max(len(columns1), len(columns2)), 4)
- grid[0, 0] = table1_label
- grid[0, 2] = table2_label
-
- for i, selector in enumerate(table1_selectors):
- grid[i + 1, 0] = widgets.Label(columns1[i])
- grid[i + 1, 1] = selector
-
- for i, selector in enumerate(table2_selectors):
- grid[i + 1, 2] = widgets.Label(columns2[i])
- grid[i + 1, 3] = selector
-
- grid[-1, :] = submit_button
- return grid
-
-
-def recommend_testing(basic_description, table_name):
- template = f"""{basic_description}
-Propose domain specific testing for table columns.
-E.g., If the table has "volume 1", "volume 2" and "total volume", then the domain specific testing should be "volume 1 + volume 2 = total volume".
-If the table has "start date" and "end date", then the domain specific testing should be "start date < end date" when these columns are not null.
-Don't write tests that checks the column range/domains, as they are already given.
-
-Now respond in the following format:
-```json
-[
-{{
- "name": "Total Volume Check",
- "reasoning": "The total volume should be the sum of volume 1 and volume 2",
- "sql" (select rows that violate the rule, to be executed by duckdb): "select * from {table_name} where volume_1 + volume_2 != total_volume",
-}},
-...
-]
-```"""
-
- messages = [{"role": "user", "content": template}]
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- processed_string = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = json.loads(processed_string)
-
- def test_json_code(code):
- if not isinstance(code, list):
- raise ValueError("The provided JSON is not a list.")
-
- required_keys = {"name", "reasoning", "sql"}
- for item in code:
- if not isinstance(item, dict):
- raise ValueError("An item in the list is not a dictionary.")
-
- if not required_keys.issubset(item.keys()):
- missing_keys = required_keys - item.keys()
- raise ValueError(f"Missing keys in an item: {missing_keys}")
-
- test_json_code(json_code)
- return json_code
-
-
-
-def recommend_join_keys(source_table_description, target_table_description):
- template = f"""Below are summary of tables. The attributes are highlighted in **bold**.
-Table 1: {source_table_description}.
-
-Table 2 : {target_table_description}.
-
-Task: Recommend a list of join keys between the two tables.
-Respond in JSON format:
-```json
-[
-{{
- "reason": "Both tables contain the same concept of ...",
- "table_1_join_keys": [attribute_1, attribute_2, ...] (case sensitive),
- "table_2_join_keys": [...]
-}},
-...
-]
-```"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- summary = response['choices'][0]['message']['content']
- assistant_message = response['choices'][0]['message']
- messages.append(assistant_message)
-
- for message in messages:
- write_log(message['content'])
- write_log("---------------------")
-
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
- def verify_json_result(json_code):
- if not isinstance(json_code, list):
- raise ValueError("JSON code should be a list")
-
- for item in json_code:
- if not isinstance(item, dict):
- raise ValueError("Each item in the list should be a dictionary")
-
- required_keys = ["reason", "table_1_join_keys", "table_2_join_keys"]
- for key in required_keys:
- if key not in item:
- raise ValueError(f"Key '{key}' is missing in one of the items")
-
- if not isinstance(item["table_1_join_keys"], list) or not isinstance(item["table_2_join_keys"], list):
- raise ValueError("table_1_join_keys and table_2_join_keys should be lists")
-
- verify_json_result(json_code)
- return json_code
-
-def plot_venn_percentage(array1, array2, name1, name2):
-
- if not isinstance(array1, np.ndarray) or not isinstance(array2, np.ndarray):
- raise TypeError("Both inputs must be numpy arrays.")
-
- plt.figure(figsize=(3, 2))
-
- set1 = set(array1)
- set2 = set(array2)
-
- total_elements = len(set1.union(set2))
-
- if total_elements == 0:
- overlap = 0
- only_set1 = 0
- only_set2 = 0
- else:
- overlap = len(set1.intersection(set2)) / total_elements * 100
- only_set1 = len(set1 - set2) / total_elements * 100
- only_set2 = len(set2 - set1) / total_elements * 100
-
- font_size=8
-
- venn_diagram = venn2(subsets=(only_set1, only_set2, overlap), set_labels=(name1, name2))
-
- for patch in venn_diagram.patches:
- if patch:
- patch.set_alpha(0.5)
-
- venn_diagram.get_label_by_id('10').set_text(f'{round(only_set1, 1)}%')
- venn_diagram.get_label_by_id('10').set_fontsize(font_size)
- venn_diagram.get_label_by_id('01').set_text(f'{round(only_set2, 1)}%')
- venn_diagram.get_label_by_id('01').set_fontsize(font_size)
- if overlap > 0 and venn_diagram.get_label_by_id('11'):
- venn_diagram.get_label_by_id('11').set_text(f'{round(overlap, 1)}%')
- venn_diagram.get_label_by_id('11').set_fontsize(font_size)
-
- for label in venn_diagram.set_labels:
- label.set_fontsize(font_size)
-
- plt.show()
-
- return overlap, only_set1, only_set2
-
-
-
-
-
-def create_column_selector(columns, callback, default=False):
- multi_select = widgets.SelectMultiple(
- options=[(column, i) for i, column in enumerate(columns)],
- disabled=False,
- layout={'width': '600px', 'height': '200px'}
- )
-
- instructions_text = "Tip: Hold Ctrl (or Cmd on Mac) to select multiple options. Currently, 0 are selected."
- instructions = widgets.Label(value=instructions_text)
-
- def update_instructions(change):
- selected_count = len(multi_select.value)
- instructions.value = f"Tip: Hold Ctrl (or Cmd on Mac) to select multiple options. Currently, {selected_count} are selected."
-
- multi_select.observe(update_instructions, 'value')
-
- select_all_button = widgets.Button(description="Select All", button_style='info', icon='check-square')
- deselect_all_button = widgets.Button(description="Deselect All", button_style='danger', icon='square-o')
- reverse_selection_button = widgets.Button(description="Reverse Selection", button_style='warning', icon='exchange')
- submit_button = widgets.Button(description="Submit", button_style='success', icon='check')
-
- def select_all(b):
- multi_select.value = tuple(range(len(columns)))
-
- def deselect_all(b):
- multi_select.value = ()
-
- def reverse_selection(b):
- current_selection = set(multi_select.value)
- all_indices = set(range(len(columns)))
- new_selection = tuple(all_indices - current_selection)
- multi_select.value = new_selection
-
- def submit(b):
- selected_indices = multi_select.value
- callback(selected_indices)
-
- select_all_button.on_click(select_all)
- deselect_all_button.on_click(deselect_all)
- reverse_selection_button.on_click(reverse_selection)
- submit_button.on_click(submit)
-
- buttons = widgets.HBox([select_all_button, deselect_all_button, reverse_selection_button])
- ui = widgets.VBox([instructions, multi_select, buttons, submit_button])
- display(ui)
-
- if default:
- multi_select.value = tuple(range(len(columns)))
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-def find_duplicate_indices(df):
- if not isinstance(df, pd.DataFrame):
- raise ValueError("Input must be a pandas DataFrame")
-
- df_temp = df.fillna("NaN_placeholder")
-
- duplicates = df_temp[df_temp.duplicated(keep=False)]
-
- if duplicates.empty:
- return []
-
- grouped = duplicates.groupby(list(duplicates.columns)).apply(lambda x: x.index.tolist())
-
- return grouped.values.tolist()
-
-
-
-
-def display_duplicated_rows_html(df, duplicated_indices):
- html_output = f"
🤨 There are {len(duplicated_indices)} groups of duplicated rows.
"
- for i, group in enumerate(duplicated_indices[:5]):
- html_output += f"Group {i+1} appear {len(group)} times:"
- for idx in group[:1]:
- html_output += df.iloc[[idx]].to_html()
- if len(duplicated_indices) > 5:
- html_output += "
...
"
- html_output += "
🧐 Do you want to remove the duplicated rows?
"
- display(HTML(html_output))
-
-
-def create_progress_bar_with_numbers(current, labels):
-
- total = len(labels)
- circles_with_labels = []
- for i in range(total):
- color = "#274e13" if i == current else "#d9ead3"
- circle_html = f'''
-
- {i+1}
-
-
- '''
- circles_with_labels.append(circle_html)
-
- display(HTML(''.join(circles_with_labels)))
-
-
-def color_columns(df, color, column_indices):
- return color_columns_multiple(df, [color], [column_indices])
-
-
-
-def columns_with_all_missing_values(df):
- if not isinstance(df, pd.DataFrame):
- raise ValueError("Input must be a pandas DataFrame")
-
- missing_column_indices = [i for i, col in enumerate(df.columns) if df.iloc[:, i].isna().all()]
-
- return missing_column_indices
-
-def display_and_ask_removal(df, missing_column_indices):
- html_output = f"
🤔 There are {len(missing_column_indices)} columns with all missing values:
"
-
- html_output += ""
- return html_output
-
-
- def generate_html_unusual_values_report(self):
-
- html_output = self.generate_unusual_values_report()
- display(HTML(html_output))
-
- def plot_missing_values_compact(self):
- """
- This function takes a pandas DataFrame as input and plots a compact bar chart showing
- the percentage of missing values for each column that has missing values.
- """
- df = self.df
-
- missing_percent = df.isnull().mean() * 100
-
- missing_percent = missing_percent[missing_percent > 0]
-
- if len(missing_percent) == 0:
- return
-
- plot_width = max(4, len(missing_percent) * 0.5)
-
- sns.set(style="whitegrid")
-
- plt.figure(figsize=(plot_width, 3), dpi=100)
- bar_plot = sns.barplot(x=missing_percent.index, y=missing_percent, color="lightgreen")
-
- for index, value in enumerate(missing_percent):
- bar_plot.text(index, value, f'{value:.2f}%', color='black', ha="center", fontsize=8)
-
- plt.title('Missing Values %', fontsize=8)
- plt.ylabel('%', fontsize=10)
- plt.xticks(rotation=45, fontsize=6)
- plt.yticks(fontsize=8)
- plt.tight_layout()
- plt.show()
-
- def get_column_warnings(self, col):
- warnings = []
- if "consistency" in self.document:
- if col in self.document["consistency"]:
- summary = self.document["consistency"][col]["summary"]
- if summary["Inconsitencies"] :
- warnings.append({"type": "Value Consistency",
- "explanation": "There are inconsistent values: " + summary["Examples"],
- "solution": ["Proceed with the transformation as is",
- "(not implemented) Clean the inconsistent values",]})
-
- if "missing_value" in self.document:
- if col in self.document["missing_value"]:
- if self.document["missing_value"][col]:
- warnings.append({"type": "Missing Value",
- "explanation": "There are missing values: " + self.document["missing_value"][col]["summary"][0],
- "solution": ["Proceed with the transformation as is",
- "Remove the rows with missing values",
- "(not implemented) Clean the inconsistent values",]})
-
- if "unusual" in self.document:
- if col in self.document["unusual"]:
- summary = self.document["unusual"][col]["summary"]
- if summary["Unusualness"]:
- warnings.append({"type": "Unusual Value",
- "explanation": "There are unusual values: " + summary["Examples"],
- "solution": ["Proceed with the transformation as is",
- "(not implemented) Clean the unusual values",]})
-
- return warnings
-
-
- def get_table_name(self):
- if self.table_name is None:
- return "table"
- else:
- if not self.table_name[0].isalpha() and self.table_name[0] != "_":
- table_name = "_" + self.table_name
- else:
- table_name = self.table_name
-
- return sanitize_table_name(table_name)
-
- def get_sample_text(self, sample_cols=None, sample_size=2, col_size=None):
- if sample_cols is None:
- sample_cols = self.df.columns
- table_name = self.get_table_name()
- return describe_df_in_natural_language(self.df[sample_cols], table_name, sample_size, num_cols_to_show=col_size)
-
- def get_basic_description(self, sample_size=2, cols = None, sample_cols=None):
-
- if cols is None:
- cols = self.df.columns
-
- table_sample = self.get_sample_text(sample_cols=sample_cols, sample_size=sample_size)
-
- column_desc = ""
-
- idx=1
- for col in cols:
- column_desc += f"{idx}. " + describe_column(self.stats, col) + "\n"
- idx += 1
-
- result = table_sample
- if len(cols) > 0:
- result += "\n\nColumn details:\n" + column_desc
-
- return result
-
- def show_progress(self, max_value):
- progress = widgets.IntProgress(
- value=1,
- min=0,
- max=max_value+1,
- step=1,
- description='',
- bar_style='',
- orientation='horizontal'
- )
-
- display(progress)
- return progress
-
- def display_tree(self, data):
-
- if self.display_html:
- html_content_updated, _, height = get_tree_html(data)
-
- display_html_iframe(html_content_updated, height=f"{height}px")
- else:
-
- import plotly.graph_objects as go
-
- def extract_hierarchy(data, parent_name='', hierarchy=None):
- if hierarchy is None:
- hierarchy = {'names': [], 'parents': []}
-
- for name, children in data.items():
- hierarchy['names'].append(name)
- hierarchy['parents'].append(parent_name)
- if isinstance(children, dict):
- extract_hierarchy(children, parent_name=name, hierarchy=hierarchy)
- elif isinstance(children, list):
- for child in children:
- hierarchy['names'].append(child)
- hierarchy['parents'].append(name)
- return hierarchy
-
- hierarchy = extract_hierarchy(data)
-
- ids = list(range(1, len(hierarchy['names']) + 1))
- name_to_id = {name: id for id, name in zip(ids, hierarchy['names'])}
- parent_ids = [name_to_id[parent] if parent in name_to_id else '' for parent in hierarchy['parents']]
-
- fig = go.Figure(go.Treemap(
- ids=ids,
- labels=hierarchy['names'],
- parents=parent_ids
- ))
-
- fig.update_layout(width=600, height=400, margin = dict(t=0, l=0, r=0, b=0))
-
- fig.show()
-
- def get_table_summary(self, overwrite=False, once=False):
- next_step = self.get_column_grouping
- if "table_summary" not in self.document:
- self.document["table_summary"] = {}
- else:
- if self.document["table_summary"] and not overwrite:
- write_log("Warning: table_summary already exists in the document.")
- if not once:
- next_step()
- return
- print("📝 Generating table summary based on renamed attributes...")
-
- progress = self.show_progress(1)
-
- main_entity = self.document["main_entity"]["summary"]
- table_sample = self.get_sample_text()
-
- max_trials = 3
-
- while max_trials > 0:
- try:
- summary, messages = get_table_summary(main_entity, table_sample)
- progress.value += 1
-
- for message in messages:
- write_log(message['content'])
- write_log("-----------------------------------")
-
- table_columns = self.df.columns.to_list()
-
-
-
- def extract_words_in_asterisks(text):
- import re
-
- pattern = r'\*\*(.*?)\*\*'
-
- matches = re.findall(pattern, text)
-
- return matches
-
- def check_sets_equality(list1, list2):
- if not set(list1) == set(list2):
- raise ValueError(f"""The two lists do not have the same set of elements.
- List 1: {sorted(list1)}
- List 2: {sorted(list2)}""")
-
- def find_extra_columns(text_columns, table_columns):
- text_columns_set = set(text_columns)
- table_columns_set = set(table_columns)
-
- if text_columns_set.issuperset(table_columns_set):
- extra_columns = text_columns_set - table_columns_set
- return list(extra_columns)
- else:
- raise ValueError(f"text_columns is not a superset of table_columns. text_columns: {text_columns}, table_columns: {table_columns}")
-
- def update_summary(summary, extra_columns):
- for column in extra_columns:
- summary = summary.replace(f"**{column}**", column)
-
- return summary
-
- text_columns = extract_words_in_asterisks(summary)
- extra_columns = find_extra_columns(text_columns, table_columns)
-
- if len(extra_columns) > 0:
- updated_summary = update_summary(summary, extra_columns)
- write_log(f"Warning: The following columns are not in the table: {extra_columns}")
- write_log(f"Updated summary: {updated_summary}")
- write_log("-----------------------------------")
- summary = updated_summary
-
-
-
- break
-
- except Exception as e:
- write_log(f"Error: {e}")
- write_log("-----------------------------------")
- max_trials -= 1
- if max_trials == 0:
- raise e
- break
-
-
- self.document["table_summary"]["summary"] = summary
- if LOG_MESSAGE_HISTORY:
- self.document["table_summary"]["history"] = messages
-
-
- html = f"Table Summary {replace_asterisks_with_tags(summary)}"
- display(HTML(html))
-
- def on_button_clicked(b):
- clear_output(wait=True)
- print("Submission received.")
- next_step()
-
- submit_button = widgets.Button(
- description='Submit',
- disabled=False,
- button_style='',
- tooltip='Click to submit',
- icon='check'
- )
-
- submit_button.on_click(on_button_clicked)
-
- display(submit_button)
-
- if self.viewer:
- on_button_clicked(submit_button)
-
-
- def get_visualization(self, overwrite=False, once=False):
- next_step = self.check_missing_all
- if "visualization" not in self.document:
- self.document["visualization"] = {}
- else:
- if self.document["visualization"] and not overwrite:
- write_log("Warning: visualization already exists in the document.")
- if not once:
- next_step()
- return
- create_progress_bar_with_numbers(1, doc_steps)
- print("📊 Generating visualization...")
-
- progress = self.show_progress(1)
-
- table_summary = self.document["table_summary"]["summary"]
-
- max_trials = 3
-
- while max_trials > 0:
- try:
- json_var, messages = generate_visualization_recommendation(table_summary)
-
- for message in messages:
- write_log(message['content'])
- write_log("-----------------------------------")
-
- def verify_json(json_var, columns):
- to_remove_indices = []
-
- for index, item in enumerate(json_var):
- if item["name"] not in ["Histogram", "Map"]:
- to_remove_indices.append(index)
- continue
-
- if item["name"] == "Histogram":
- if item["params"] not in columns:
- to_remove_indices.append(index)
- elif item["name"] == "Map":
- if any(param not in columns for param in item["params"]):
- to_remove_indices.append(index)
-
- for index in reversed(to_remove_indices):
- del json_var[index]
-
- return json_var
-
- json_var = verify_json(json_var, self.df.columns)
-
- break
-
- except Exception as e:
- write_log(f"Error: {e}")
- write_log("-----------------------------------")
- max_trials -= 1
- if max_trials == 0:
- raise e
- continue
-
- progress.value += 1
-
- self.document["visualization"]["summary"] = json_var
- if LOG_MESSAGE_HISTORY:
- self.document["visualization"]["history"] = messages
-
- html = self.generate_visualizations_html()
-
- display(HTML(html))
-
- def on_button_clicked(b):
- clear_output(wait=True)
- print("Submission received.")
- next_step()
-
- submit_button = widgets.Button(
- description='Submit',
- disabled=False,
- button_style='',
- tooltip='Click to submit',
- icon='check'
- )
-
- submit_button.on_click(on_button_clicked)
-
- display(submit_button)
-
- if self.viewer:
- on_button_clicked(submit_button)
-
-
- def get_column_grouping(self, overwrite=False, once=False):
- next_step = self.get_visualization
- if "column_grouping" not in self.document:
- self.document["column_grouping"] = {}
- else:
- if self.document["column_grouping"] and not overwrite:
- write_log("Warning: column_grouping already exists in the document.")
- if not once:
- next_step()
- return
- create_progress_bar_with_numbers(1, doc_steps)
- print("🌳 Building concept map...")
-
- progress = self.show_progress(1)
-
- meanings = self.document["column_meaning"]["summary"]
-
- column_meanings = '\n'.join([f"- {m['column']}: {m['meaning']}" for m in meanings])
-
- main_entity = self.document["main_entity"]["summary"]
-
- sample = self.get_sample_text()
-
-
-
-
-
- template = f"""{sample}
-
-{column_meanings}
-
-This table is about {main_entity}. The goal is to build a mind map.
-
-Recursively group the attributes based on inherent entity association, not conceptual similarity.
- E.g., for [student name, student grade, teacher grade], group by student and teacher, not by name.
-Avoid groups with too many/few subgroups.
-
-Conclude with the final result as a multi-level JSON. Make sure all attributes are included.
-
-```json
-{{
- "{main_entity}":
- {{
- "Sub group": {{
- "sub-sub group": ["attribute1", "attribute2", ...],
- }},
- }}
-}}
-```"""
-
- def extract_attributes(json_var):
- attributes = []
-
- def traverse(element):
- if isinstance(element, dict):
- for value in element.values():
- traverse(value)
-
- elif isinstance(element, list):
- for item in element:
- if isinstance(item, str):
- attributes.append(item)
-
- else:
- traverse(item)
-
-
- traverse(json_var)
-
- return attributes
-
- def validate_attributes(attributes, reference_attributes):
- error_messages = []
-
- seen_attributes = set()
- duplicates = set()
- for attribute in attributes:
- if attribute in seen_attributes:
- duplicates.add(attribute)
- seen_attributes.add(attribute)
-
- if duplicates:
- error_messages.append("Duplicate attributes: " + ', '.join(duplicates))
-
- attributes_set = set(attributes)
- reference_set = set(reference_attributes)
-
- extra_attributes = attributes_set - reference_set
- if extra_attributes:
- error_messages.append("Extra attributes: " + ', '.join(extra_attributes))
-
- missing_attributes = reference_set - attributes_set
- if missing_attributes:
- error_messages.append("Missing attributes: " + ', '.join(missing_attributes) + "\n Are attributes in the leaf as an array [att1, att2]?")
-
- return '\n'.join(error_messages)
-
-
-
- def build_concept_map_and_verify(messages):
-
- number_of_trials = 3
-
- for messgae in messages:
- write_log(messgae['content'])
- write_log("-----------------------------------")
-
- while number_of_trials > 0:
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(response['choices'][0]['message']['content'])
- write_log("-----------------------------------")
-
- assistant_message = response['choices'][0]['message']
- json_code = extract_json_code_safe(assistant_message['content'])
- json_code = json_code.replace('\'', '\"')
- json_var = json.loads(json_code)
- attributes = extract_attributes(json_var)
-
- messages.append(assistant_message)
-
- error = validate_attributes(attributes, self.df.columns)
-
- if error!= '':
- error_message = {
- "role": "user",
- "content": f"{error}\nPlease correct your answer and return the json in the required format."
- }
- messages.append(error_message)
- write_log(error_message['content'])
- write_log("-----------------------------------")
-
- number_of_trials -= 1
-
- if number_of_trials == 0:
- raise Exception(error)
-
- else:
- self.document["column_grouping"]["summary"] = json_var
- if LOG_MESSAGE_HISTORY:
- if "history" not in self.document["column_grouping"]:
- self.document["column_grouping"]["history"] = messages
- else:
- self.document["column_grouping"]["history"] += messages
- break
-
-
- messages =[ {"role": "user", "content": template}]
- build_concept_map_and_verify(messages)
- data = self.document["column_grouping"]["summary"]
-
-
-
- progress.value += 1
-
- clear_output(wait=True)
-
- self.display_tree(data)
-
- def create_widgets_for_column_grouping():
- def on_value_change(change):
- if change['new'] == 'No':
- feedback_container.layout.display = ''
- text_area.disabled = False
- else:
- feedback_container.layout.display = 'none'
- text_area.disabled = True
-
- accuracy_question_label = widgets.Label(value='Is the mind map accurate?')
-
- accuracy_check = widgets.RadioButtons(
- options=['Yes', 'No'],
- description='',
- disabled=False
- )
-
- label = widgets.Label(value='If not accurate, how to fix it?')
-
- text_area = widgets.Textarea(
- value='',
- placeholder='Type here',
- description='',
- disabled=True
- )
-
- feedback_container = widgets.VBox([label, text_area], layout=Layout(display='none'))
-
- submit_button = widgets.Button(
- description='Submit',
- disabled=False,
- button_style='',
- tooltip='Click to submit',
- icon='check'
- )
-
- accuracy_check.observe(on_value_change, names='value')
-
- def on_button_clicked(b):
- if accuracy_check.value == 'No':
- if text_area.value == '':
- print("\033[91mPlease enter the information\033[0m.")
- return
- feedback = text_area.value
-
- print("🌳 Refining concept map...")
- progress = self.show_progress(1)
-
- data = self.document["column_grouping"]["summary"]
-
- messages =[ {"role": "user", "content": template}]
- messages.append({"role": "system", "content": "```json\n"+str(data)+"\n```"})
- messages.append({"role": "user",
- "content": f"""{feedback} Please refine the json and return the result within ```json``` block."""})
-
- build_concept_map_and_verify(messages)
- clear_output(wait=True)
- data = self.document["column_grouping"]["summary"]
- self.display_tree(data)
- accuracy_question = create_widgets_for_column_grouping()
- display(accuracy_question)
-
- else:
- clear_output(wait=True)
- print("Submission received.")
- next_step()
-
- submit_button.on_click(on_button_clicked)
-
- accuracy_question = widgets.VBox([accuracy_question_label, accuracy_check, feedback_container, submit_button])
-
- return accuracy_question, on_button_clicked
-
- accuracy_question, on_button_clicked = create_widgets_for_column_grouping()
-
- display(accuracy_question)
-
- if self.viewer:
- on_button_clicked(None)
-
-
-
- def execute_project(self):
- next_step = self.check_duplicated_rows
-
- if self.project_step_finished:
- next_step()
- return
-
- keep_column_indices = self.document["project"]
- remove_column_indices = [i for i in range(len(self.df.columns)) if i not in keep_column_indices]
-
- if len(remove_column_indices) > 0:
- remove_columns_step = RemoveColumnsStep(sample_df = self.df[:2], col_indices = remove_column_indices, name="Project Columns")
- self.pipeline.add_step_to_final(remove_columns_step)
- self.df = self.pipeline.run_codes()
-
- self.project_step_finished = True
- next_step()
-
- def execute_drop_x_y_columns(self):
- next_step = self.check_index_columns
-
- if self.drop_x_y_columns_step_finished:
- next_step()
- return
-
- remove_column_indices = self.document["x_y_columns"]["remove_columns"]
-
- if len(remove_column_indices) > 0:
- remove_columns_step = RemoveColumnsStep(sample_df = self.df[:2], col_indices = remove_column_indices, name="Remove _x _y Columns")
- self.pipeline.add_step_to_final(remove_columns_step)
- self.df = self.pipeline.run_codes()
-
- self.drop_x_y_columns_step_finished = True
- next_step()
-
- def execute_drop_duplicated_columns(self):
- next_step = self.check_x_y_columns
-
- if self.drop_duplicated_colulmn_step_finished:
- next_step()
- return
-
- remove_column_indices = self.document["duplicate_columns"]["remove_columns"]
-
- if len(remove_column_indices) > 0:
- remove_columns_step = RemoveColumnsStep(sample_df = self.df[:2], col_indices = remove_column_indices, name="Remove Duplicated Columns")
- self.pipeline.add_step_to_final(remove_columns_step)
- self.df = self.pipeline.run_codes()
-
- self.drop_duplicated_colulmn_step_finished = True
- next_step()
-
- def execute_index_columns(self):
- next_step = self.check_data_type
-
- if self.drop_index_column_step_finished:
- next_step()
- return
-
- index_column_indices = self.document["index_columns"]["remove_columns"]
-
- if len(index_column_indices) > 0:
- remove_columns_step = RemoveColumnsStep(sample_df = self.df[:2], col_indices = index_column_indices, name="Remove Index Columns")
- self.pipeline.add_step_to_final(remove_columns_step)
- self.df = self.pipeline.run_codes()
-
- self.drop_index_column_step_finished = True
- next_step()
-
- def execute_drop_all_missing_columns(self):
- next_step = self.check_duplicated_columns
-
- if self.drop_all_missing_columns_step_finished:
- next_step()
- return
-
- remove_column_indices = self.document["missing_columns"]["remove_columns"]
-
- if len(remove_column_indices) > 0:
- remove_columns_step = RemoveColumnsStep(sample_df = self.df[:2], col_indices = remove_column_indices, name="Remove All Missing Columns")
- self.pipeline.add_step_to_final(remove_columns_step)
- self.df = self.pipeline.run_codes()
-
- self.drop_all_missing_columns_step_finished = True
- next_step()
-
- def execute_deduplicated_rows(self):
- next_step = self.check_all_missing_columns
-
- if self.drop_duplicated_row_step_finished:
- next_step()
- return
-
- if self.document["duplicate_rows"]["remove_duplicates"]:
- deduplicate_step = RemoveDuplicatesStep(sample_df = self.df[:2])
- self.pipeline.add_step_to_final(deduplicate_step)
- self.df = self.pipeline.run_codes()
-
- self.drop_duplicated_row_step_finished = True
- next_step()
-
- def execute_remove_data_type(self):
- next_step = self.get_main_entity
-
- if self.remove_data_type_step_finished:
- next_step()
- return
-
- column_data_type_dict = {}
-
- for column_name in self.document["data_type"]:
- data_type = self.document["data_type"][column_name]["data_type"]
- if "invalid_rows" in self.document["data_type"][column_name]:
- column_data_type_dict[column_name] = data_type
-
- if column_data_type_dict != {}:
- clean_data_type_step = CleanDataType(sample_df = self.df[:2], column_data_type_dict = column_data_type_dict)
- self.pipeline.add_step_to_final(clean_data_type_step)
- self.df = self.pipeline.run_codes()
-
- self.remove_data_type_step_finished = True
- next_step()
-
-
- def execute_rename_column(self):
-
- next_step = self.get_table_summary
-
- if self.rename_step_finished:
- self.stats = collect_df_statistics(self)
- next_step()
- return
-
- rename_summary = self.document["rename_column"]["summary"]
-
- self.name_mapping = {}
-
- if self.rename_table:
- for item in rename_summary:
- if item['rename'] != '':
- print(f"Renaming column {item['column']} to {item['rename']}")
- self.name_mapping[item['column']] = item['rename']
-
- if not self.name_mapping == {}:
- rename_step = ColumnRename(sample_df = self.df[:2], rename_map = self.name_mapping)
- self.pipeline.add_step_to_final(rename_step)
- self.df = self.pipeline.run_codes()
-
- self.rename_step_finished = True
-
-
-
-
-
- self.stats = collect_df_statistics(self)
-
-
- next_step()
-
- def generate_pipeline(self):
- return self.pipeline
-
- def decide_project(self, overwrite=False, once=False):
-
- next_step = self.execute_project
-
- if "project" not in self.document:
- pass
- else:
- if not overwrite:
- write_log("Warning: project already exists in the document.")
- if not once:
- next_step()
- return
-
- create_progress_bar_with_numbers(0, doc_steps)
-
- df = self.df
-
- df_sample = df[:100]
- display(HTML(wrap_in_scrollable_div(truncate_html_td(df_sample.to_html()))))
-
- num_cols = len(df.columns)
-
- display(HTML(f"
🧐 There are {num_cols} columns. Please select the columns that you want to keep.
"))
-
- column_names = self.df.columns.to_list()
-
- def callback_next(selected_indices):
- clear_output(wait=True)
- self.document["project"] = selected_indices
- next_step()
-
- create_column_selector(column_names, callback_next, default=True)
-
- if self.viewer:
- callback_next(list(range(num_cols)))
-
-
- def check_duplicated_rows(self, overwrite=False, once=False):
-
- next_step = self.execute_deduplicated_rows
-
- if "duplicate_rows" not in self.document or \
- "duplicated_indices" not in self.document["duplicate_rows"] or\
- "remove_duplicates" not in self.document["duplicate_rows"]:
- self.document["duplicate_rows"] = {}
- else:
- if self.document["duplicate_rows"] and not overwrite:
- write_log("Warning: duplicate_rows already exists in the document.")
- if not once:
- next_step()
- return
-
- create_progress_bar_with_numbers(0, doc_steps)
- print("🔍 Checking duplicated rows...")
-
- duplicated_indices = find_duplicate_indices(self.df)
- self.document["duplicate_rows"]["duplicated_indices"] = duplicated_indices
- self.document["duplicate_rows"]["num_duplicated_rows"] = len(duplicated_indices)
-
- if len(duplicated_indices) > 0:
- display_duplicated_rows_html(self.df, duplicated_indices)
-
- def on_button_clicked(b):
- clear_output(wait=True)
- if b.description == 'Yes':
- self.document["duplicate_rows"]["remove_duplicates"] = True
- else:
- self.document["duplicate_rows"]["remove_duplicates"] = False
- next_step()
-
- yes_button = widgets.Button(
- description='Yes',
- disabled=False,
- button_style='',
- tooltip='Click to submit Yes',
- icon='check'
- )
-
- no_button = widgets.Button(
- description='No',
- disabled=False,
- button_style='',
- tooltip='Click to submit No',
- icon='times'
- )
-
- yes_button.on_click(on_button_clicked)
- no_button.on_click(on_button_clicked)
-
- display(HBox([yes_button, no_button]))
-
- if self.viewer:
- on_button_clicked(yes_button)
-
-
- else:
- self.document["duplicate_rows"]["remove_duplicates"] = False
- next_step()
-
- def check_index_columns(self, overwrite=False, once=False):
- next_step = self.execute_index_columns
-
- if "index_columns" not in self.document or\
- "index_column_indices" not in self.document["index_columns"] or \
- "remove_columns" not in self.document["index_columns"]:
- self.document["index_columns"] = {}
- else:
- if self.document["index_columns"] and not overwrite:
- write_log("Warning: index_columns already exists in the document.")
- if not once:
- next_step()
- return
-
- clear_output(wait=True)
- create_progress_bar_with_numbers(0, doc_steps)
- print("🔍 Checking index columns...")
-
- index_column_indices = find_default_index_column(self.df)
-
- self.document["index_columns"]["index_column_indices"] = index_column_indices
- self.document["index_columns"]["index_column_names"] = [self.df.columns[i] for i in index_column_indices]
-
- if len(index_column_indices) > 0:
- display_index_and_ask_removal(self.df, index_column_indices)
-
- def callback_next(selected_indices):
- to_remove_indices = [index_column_indices[i] for i in selected_indices]
-
- self.document["index_columns"]["remove_columns"] = to_remove_indices
- next_step()
-
- create_column_selector(self.df.columns.to_list(), callback_next)
-
- if self.viewer:
- callback_next([])
-
- else:
- self.document["index_columns"]["remove_columns"] = []
- next_step()
-
-
-
- def check_all_missing_columns(self, overwrite=False, once=False):
-
- next_step = self.execute_drop_all_missing_columns
-
- if "missing_columns" not in self.document or\
- "missing_column_indices" not in self.document["missing_columns"] or\
- "remove_columns" not in self.document["missing_columns"]:
- self.document["missing_columns"] = {}
- else:
- if self.document["missing_columns"] and not overwrite:
- write_log("Warning: missing_columns already exists in the document.")
- if not once:
- next_step()
- return
-
- clear_output(wait=True)
- create_progress_bar_with_numbers(0, doc_steps)
- print("🔍 Checking columns with all missing values...")
-
- missing_column_indices = columns_with_all_missing_values(self.df)
- self.document["missing_columns"]["missing_column_indices"] = missing_column_indices
- self.document["missing_columns"]["missing_column_names"] = [self.df.columns[i] for i in missing_column_indices]
-
- if len(missing_column_indices) > 0:
- display_and_ask_removal(self.df, missing_column_indices)
-
- column_names = self.df.columns.to_list()
-
- missing_columns = [column_names[i] for i in missing_column_indices]
-
- def callback_next(selected_indices):
- to_remove_indices = [missing_column_indices[i] for i in selected_indices]
- self.document["missing_columns"]["remove_columns"] = to_remove_indices
- next_step()
-
- create_column_selector(missing_columns, callback_next)
-
- if self.viewer:
- callback_next(list(range(len(missing_columns))))
-
- else:
- self.document["missing_columns"]["remove_columns"] = []
- next_step()
-
- def check_duplicated_columns(self, overwrite=False, once=False):
- next_step = self.execute_drop_duplicated_columns
-
- if "duplicate_columns" not in self.document or\
- "duplicated_column_indices" not in self.document["duplicate_columns"] or\
- "remove_columns" not in self.document["duplicate_columns"]:
- self.document["duplicate_columns"] = {}
- else:
- if self.document["duplicate_columns"] and not overwrite:
- write_log("Warning: duplicate_columns already exists in the document.")
- if not once:
- next_step()
- return
-
- clear_output(wait=True)
- create_progress_bar_with_numbers(0, doc_steps)
- print("🔍 Checking duplicated columns...")
-
- duplicated_column_indices = find_duplicate_column_indices(self.df)
-
- self.document["duplicate_columns"]["duplicated_column_indices"] = duplicated_column_indices
- self.document["duplicate_columns"]["duplicated_column_names"] = [[self.df.columns[i] for i in group_indices] for group_indices in duplicated_column_indices]
-
- if len(duplicated_column_indices) > 0:
- display_duplicated_columns_html(self.df, duplicated_column_indices)
-
- column_names = self.df.columns.to_list()
-
- duplicated_columns = []
- for group_indices in duplicated_column_indices:
- duplicated_columns += [f"{column_names[i]} ({i+1})" for i in group_indices]
-
- def callback_next(selected_indices):
-
- flattened_indices = [item for sublist in duplicated_column_indices for item in sublist]
-
- to_remove_indices = [flattened_indices[i] for i in selected_indices]
-
- self.document["duplicate_columns"]["remove_columns"] = to_remove_indices
- next_step()
-
- create_column_selector(duplicated_columns, callback_next)
-
- if self.viewer:
- callback_next(list(range(len(duplicated_columns))))
-
- else:
- self.document["duplicate_columns"]["remove_columns"] = []
- next_step()
-
- def check_x_y_columns(self, overwrite=False, once=False):
- next_step = self.execute_drop_x_y_columns
-
- if "x_y_columns" not in self.document or \
- "x_y_column_indices" not in self.document["x_y_columns"] or \
- "remove_columns" not in self.document["x_y_columns"]:
- self.document["x_y_columns"] = {}
- else:
- if self.document["x_y_columns"] and not overwrite:
- write_log("Warning: x_y_columns already exists in the document.")
- if not once:
- next_step()
- return
-
- clear_output(wait=True)
- create_progress_bar_with_numbers(0, doc_steps)
-
- print("🔍 Checking x and y columns...")
-
- x_y_column_indices = find_columns_with_xy_suffix_indices(self.df)
-
- self.document["x_y_columns"]["x_y_column_indices"] = x_y_column_indices
- self.document["x_y_columns"]["x_y_column_names"] = [[self.df.columns[i] for i in group_indices] for group_indices in x_y_column_indices]
-
- if len(x_y_column_indices) > 0:
- display_xy_duplicated_columns_html(self.df, x_y_column_indices)
-
- column_names = self.df.columns.to_list()
-
- x_y_columns = []
-
- for group_indices in x_y_column_indices:
- x_y_columns += [column_names[i] for i in group_indices]
-
- def callback_next(selected_indices):
-
- flattened_indices = [item for sublist in x_y_column_indices for item in sublist]
-
- to_remove_indices = [flattened_indices[i] for i in selected_indices]
-
- self.document["x_y_columns"]["remove_columns"] = to_remove_indices
- next_step()
-
- create_column_selector(x_y_columns, callback_next)
-
- if self.viewer:
- callback_next(list(range(len(x_y_columns))))
-
- else:
- self.document["x_y_columns"]["remove_columns"] = []
- next_step()
-
- def check_data_type(self, overwrite=False, once=False):
- next_step = self.execute_remove_data_type
-
- if "data_type" in self.document:
- if self.document["data_type"] and not overwrite:
- write_log("Warning: data_type already exists in the document.")
- if not once:
- next_step()
- return
-
-
- clear_output(wait=True)
-
- clear_output(wait=True)
- create_progress_bar_with_numbers(0, doc_steps)
- print("🔍 Checking data types...")
-
- table_name = self.get_table_name()
-
- duckdb_conn = duckdb.connect()
- duckdb_conn.register(table_name, self.df)
-
- schema_query = f"PRAGMA table_info('{table_name}')"
- schema_info = duckdb_conn.execute(schema_query).df()
-
- self.document["data_type"] = {}
-
- has_invalid_data_type = False
-
- df = self.df
-
- for index, row in schema_info.iterrows():
- column_name = row['name']
- data_type = row['type'].upper()
-
-
-
- mask = select_invalid_data_type(df, column_name, data_type)
-
- self.document["data_type"][column_name] = {"data_type": data_type}
-
- if mask.any():
- print(f"{BOLD}Column '{column_name}'{END} is of type {ITALIC}{data_type}{END}.")
-
-
- example = df[mask][column_name][:5].to_list()
- self.document["data_type"][column_name]["invalid_rows"] = example
-
- print(f" ⚠️ There are {mask.sum()} invalid rows that don't match the data type.")
- print(f" ⚠️ Examples: {example}")
-
- has_invalid_data_type = True
-
- if has_invalid_data_type:
- print("⚠️ These invalid rows have to be removed.")
- print("😊 Support for more flexible cleaning is coming soon.")
-
- def on_button_clicked(b):
- clear_output(wait=True)
- next_step()
-
- next_button = widgets.Button(
- description='Next',
- disabled=False,
- button_style='',
- tooltip='Click to go to the next step',
- icon='check'
- )
-
- next_button.on_click(on_button_clicked)
-
- display(next_button)
-
- if self.viewer:
- on_button_clicked(next_button)
-
- else:
- next_step()
-
-
- def rename_column(self, overwrite=False, once=False):
- next_step = self.execute_rename_column
- if "rename_column" not in self.document:
- self.document["rename_column"] = {}
- else:
- if self.document["rename_column"] and not overwrite:
- write_log("Warning: rename_column already exists in the document.")
- if not once:
- next_step()
- return
- create_progress_bar_with_numbers(1, doc_steps)
- print("🏷️ Renaming the columns...")
- progress = self.show_progress(1)
-
- meanings = self.document["column_meaning"]["summary"]
-
- max_trials = 3
-
- while max_trials > 0:
- try:
- json_code, messages = get_rename(meanings)
- progress.value += 1
- break
-
- except Exception as e:
- write_log(f"Error: {e}")
- write_log("-----------------------------------")
- max_trials -= 1
- if max_trials == 0:
- raise e
- continue
-
-
-
-
-
-
-
-
- self.document["rename_column"]["summary"] = json_code
- if LOG_MESSAGE_HISTORY:
- self.document["rename_column"]["history"] = messages
-
- next_step()
-
- def get_column_meaning(self, overwrite=False, once=False):
- next_step = self.rename_column
- if "column_meaning" not in self.document:
- self.document["column_meaning"] = {}
- else:
- if self.document["column_meaning"] and not overwrite:
- write_log("Warning: column_meaning already exists in the document.")
- if not once:
- next_step()
- return
- create_progress_bar_with_numbers(1, doc_steps)
- print("💡 Understanding the columns...")
- progress = self.show_progress(1)
-
- main_entity = self.document["main_entity"]["summary"]
- basic_description = self.get_basic_description()
-
-
-
-
-
-
- template = f"""{basic_description}
-
-This table is about {main_entity}. The goal is study the high-level column meaning.
-Use short simple words to describe the most possible meanings.
-Respond in JSON, for all columns:
-```json
-[{{
- "column": "column_name",
- "meaning": "short, simple guess on the meaning"
-}},...]
-```"""
-
- max_trials = 3
-
- while max_trials > 0:
- try:
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- progress.value += 1
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- data = json.loads(json_code)
-
- messages.append(response['choices'][0]['message'])
-
- def check_column_complete(data):
- if len(data) != len(self.df.columns):
- raise Exception("Not all columns are covered in the column meaning.")
-
- for item in data:
- if item['column'] not in self.df.columns:
- raise Exception(f"Column {item['column']} does not exist in the table.")
-
- check_column_complete(data)
-
- self.document["column_meaning"]["summary"] = data
- if LOG_MESSAGE_HISTORY:
- self.document["column_meaning"]["history"] = messages
-
- break
-
- except Exception as e:
- write_log(f"Error: {e}")
- write_log("-----------------------------------")
- max_trials -= 1
- if max_trials == 0:
- raise e
- continue
-
- clear_output(wait=True)
-
- def radio_change_handler(change):
- instance = change['owner']
- if instance.value == 'Other':
- instance.text_area.layout.display = ''
- else:
- instance.text_area.layout.display = 'none'
-
- container = widgets.VBox()
-
- column_to_radio = {}
-
- for item in data:
- if 'ambiguous' in item and item['ambiguous']:
- label_text = f"{item['column']}: (This column has ambiguous interpretations)"
- else:
- label_text = f"{item['column']}:"
-
- label = widgets.HTML(
- value=label_text,
- layout=widgets.Layout(margin='0 0 10px 0')
- )
-
- options = [item['meaning']]
- if 'ambiguous' in item:
- options += item['ambiguous']
- options += ['Other']
-
- radio = widgets.RadioButtons(
- options=options,
- value=item['meaning'],
- layout=widgets.Layout(width='80%', align_items='flex-start')
- )
-
- text_area = widgets.Textarea(
- value='',
- placeholder='Please provide the meaning of the column.',
- description='',
- disabled=False,
- layout=widgets.Layout(display='none', width='100%')
- )
-
- radio.text_area = text_area
- column_to_radio[item['column']] = radio
- radio.observe(radio_change_handler, names='value')
-
- container.children += (label, radio, text_area)
-
- def submit_callback(btn):
- error_items = []
-
- for item in data:
- radio = column_to_radio[item['column']]
- if radio.value == 'Other' and not radio.text_area.value.strip():
- error_items.append(item)
-
- if error_items:
- clear_output(wait=True)
- display(container, submit_btn)
- for item in error_items:
- print(f"\033[91m{item['column']} meaning can't be empty.\033[0m")
- else:
- clear_output(wait=True)
- feedback_data = []
-
- for item in data:
- radio = column_to_radio[item['column']]
-
- if radio.value == 'Other':
- feedback_data.append({'column': item['column'], 'meaning': radio.text_area.value})
- else:
- feedback_data.append({'column': item['column'], 'meaning': radio.value})
- self.document["column_meaning"]["summary"] = feedback_data
- if LOG_MESSAGE_HISTORY:
- self.document["column_meaning"]["history"].append({"role": "user",
- "content": feedback_data})
- print("Submission received.")
- next_step()
-
-
-
- submit_btn = widgets.Button(
- description="Submit",
- button_style='',
- tooltip='Submit',
- icon=''
- )
-
- submit_btn.on_click(submit_callback)
-
- display(container, submit_btn)
-
- if self.viewer:
- submit_callback(submit_btn)
-
-
- def get_main_entity(self, overwrite=False, once=False):
-
- next_step = self.get_column_meaning
-
- if "main_entity" not in self.document:
- self.document["main_entity"] = {}
- else:
- if self.document["main_entity"] and not overwrite:
- write_log("Warning: main_entity already exists in the document.")
- if not once:
- next_step()
- return
-
- create_progress_bar_with_numbers(1, doc_steps)
-
- print("💡 Understanding the table...")
- progress = self.show_progress(1)
-
- basic_description = self.get_basic_description()
-
- template = f"""{basic_description}
-
-Identify the main entity this table is mainly recording.
-- Entity is tangible or conceptual thing that can exist or be conceptualized individually.
- Example of Entity: Person, Course, Location, Time, School, Job
-- Entity is not property, or aspect. Infer the main underlying entities.
- Example of Non-entity: Height, Weight, Information, Name, Finance
- For height, the main underlying entity is "People"
-- Entity shall be a clear single phrase. Don't use / to combine entities.
-
-1. Start by reasoning what the table is about, and the candidate main entities.
- If there are multiple candidate main entities, discuss if there is a main relationship entity that can be used to group them.
-2. Conclude by listing the main entity.
-
-Now respond in the following format:
-```json
-{{
- "reasoning": "The table is about ...",
- "main entity": "..."
-}}
-```"""
-
- messages = [{"role": "user", "content": template}]
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- progress.value += 1
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- processed_string = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = json.loads(processed_string)
- main_entity = json_code['main entity']
-
- messages.append(response['choices'][0]['message'])
- self.document["main_entity"]["summary"] = main_entity
- self.document["main_entity"]["reasoning"] = json_code['reasoning']
- if LOG_MESSAGE_HISTORY:
- self.document["main_entity"]["history"] = messages
-
- clear_output(wait=True)
- reason = json_code['reasoning']
- print(f'\033[1mThe table is mainly talking about: \033[0m{main_entity}')
- print(f'\033[1mDetails\033[0m: {reason}')
-
- accuracy_question_label = widgets.Label(value='Is the above information accurate?')
-
- accuracy_check = widgets.RadioButtons(
- options=['Yes', 'No'],
- description='',
- disabled=False
- )
-
- label = widgets.Label(value='If not accurate, the table is mainly talking about:')
-
- text_area = widgets.Textarea(
- value='',
- placeholder='Type here',
- description='',
- disabled=True
- )
-
- feedback_container = widgets.VBox([label, text_area], layout=Layout(display='none'))
-
- submit_button = widgets.Button(
- description='Submit',
- disabled=False,
- button_style='',
- tooltip='Click to submit',
- icon='check'
- )
-
- def on_button_clicked(b):
- clear_output(wait=True)
- if accuracy_check.value == 'No':
- if text_area.value == '':
- display(accuracy_question)
- print("\033[91mPlease enter the information\033[0m.")
- return
- corrected_entity = text_area.value
-
- print(f"Corrected information received. This table is mainly about {corrected_entity}")
- self.document["main_entity"]["summary"] = corrected_entity
- if LOG_MESSAGE_HISTORY:
- self.document["main_entity"]["history"].append({"role": "user",
- "content": template})
- else:
- print("Submission received.")
- next_step()
-
- def on_value_change(change):
- if change['new'] == 'No':
- feedback_container.layout.display = ''
- text_area.disabled = False
- else:
- feedback_container.layout.display = 'none'
- text_area.disabled = True
-
- accuracy_check.observe(on_value_change, names='value')
-
- submit_button.on_click(on_button_clicked)
-
- accuracy_question = widgets.VBox([accuracy_question_label, accuracy_check, feedback_container, submit_button])
-
- display(accuracy_question)
-
- if self.viewer:
- on_button_clicked(submit_button)
-
-
- def document_all(self):
- self.get_main_entity()
- self.get_column_grouping()
- self.check_consistency_all()
- self.check_pattern_all()
- self.check_missing_all()
- self.check_unusual_all()
-
- def check_missing_all(self):
- next_step = self.check_unusual_all
-
- create_progress_bar_with_numbers(2, doc_steps)
- print("❓ Checking the missing values...")
- progress = self.show_progress(len(self.df.columns))
-
- for col in self.df.columns:
- self.check_missing(col)
-
- progress.value += 1
-
- ambiguous_missing = {}
-
- for col in self.document["missing_value"]:
- if "summary" in self.document["missing_value"][col]:
- summary = self.document["missing_value"][col]["summary"]
- if isinstance(summary, list):
- ambiguous_missing[col] = summary
-
- if not ambiguous_missing:
- next_step()
- return
-
- clear_output(wait=True)
- print("The following columns have missing values: ❓")
-
- def radio_change_handler(change):
- instance = change['owner']
- if instance.value == 'Other':
- instance.text_area.layout.display = ''
- else:
- instance.text_area.layout.display = 'none'
-
- container = widgets.VBox()
-
-
- col_to_radio = {}
-
- for item in ambiguous_missing:
-
- label_text = f"{item}"
-
- label = widgets.HTML(value=label_text)
-
- reasons = ambiguous_missing[item]
- options = [f"{reason['class']}: {reason['explanation']}" for reason in reasons] + ['Unclear','Other']
-
- radio = widgets.RadioButtons(
- options=options,
- value=options[0],
- layout=widgets.Layout(width='80%', align_items='flex-start')
- )
-
- text_area = widgets.Textarea(
- value='',
- placeholder='Please provide the reason for missing values.',
- description='',
- disabled=False,
- layout=widgets.Layout(display='none', width='100%')
- )
-
- radio.text_area = text_area
- col_to_radio[item] = radio
- radio.observe(radio_change_handler, names='value')
-
- item_container = widgets.VBox([label, radio, text_area])
- container.children += (item_container,)
-
- def submit_callback(btn):
- error_items = []
-
- for col in col_to_radio:
- radio = col_to_radio[col]
- if radio.value == 'Other' and not radio.text_area.value.strip():
- error_items.append(col)
-
- if error_items:
- for col in error_items:
- print(f"\033[91m{col} reason can't be empty.\033[0m")
- else:
- clear_output(wait=True)
-
- for col in col_to_radio:
- radio = col_to_radio[col]
-
- if radio.value == 'Other':
- self.document["missing_value"][col]["summary"] = radio.text_area.value
- if LOG_MESSAGE_HISTORY:
- self.document["missing_value"][col]["history"].append({"role": "user",
- "content": radio.text_area.value})
- else:
- self.document["missing_value"][col]["summary"] = radio.value
- if LOG_MESSAGE_HISTORY:
- self.document["missing_value"][col]["history"].append({"role": "user",
- "content": radio.value})
-
- print("Submission received.")
- next_step()
-
-
-
- submit_btn = widgets.Button(
- description="Submit",
- button_style='',
- tooltip='Submit',
- icon=''
- )
-
- submit_btn.on_click(submit_callback)
-
- display(container, submit_btn)
-
- if self.viewer:
- submit_callback(submit_btn)
-
-
-
-
- def check_missing(self, column: str):
- if column not in self.df.columns:
- raise ValueError(f"Column {column} does not exist in the DataFrame.")
-
- if "missing_value" not in self.document:
- self.document["missing_value"] = {}
-
- if column not in self.document["missing_value"]:
- self.document["missing_value"][column] = {}
- else:
- if self.document["missing_value"][column]:
- write_log(f"Warning: {column} already exists in the document.")
- return
-
- df_col = self.df[column]
-
- if isinstance(df_col, pd.DataFrame):
- df_col = df_col.iloc[:, 0]
-
- nan_rows = self.df[df_col.isna()]
- non_nan_rows = self.df.dropna(subset=[column])
-
- nan_sample = nan_rows.head(3)
- non_nan_sample = non_nan_rows.head(3)
-
- if len(nan_sample) == 0:
- write_log(f"Warning: {column} does not have missing values.")
- return
-
- nan_sample_str = nan_sample.to_string(index=False)
- non_nan_sample_str = ""
- if len(non_nan_sample) > 0:
- non_nan_sample_str = f"Sample of data without missing values:\n{non_nan_sample.to_string(index=False)}"
- else:
- non_nan_sample_str = "The whole column has missing values."
-
- main_entity = self.document["main_entity"]["summary"]
-
- template = f"""In a table about {main_entity}, {column} has missing value.
-
-Sample of data with missing values:
-{nan_sample_str}
-
-{non_nan_sample_str}
-
-There are general 5 classes of missing values:
- Not Applicable: Certain questions or fields do not apply to the individual/entity being measured. For example, a question about "spouse's occupation" wouldn't apply to someone who is unmarried.
- Censorship: for sensitive attribute, the data can be masked for privacy
- Non-Response: The subject chose not to provide information or ignored the request. This is common in surveys and certain types of observational research.
- Not Collected: Information wasn't gathered due to oversight, resource limitations (e.g., certain tests or measures could not be performed), or it was deemed unnecessary at the time of collection.
- Damaged Data: Information was originally collected but later became unavailable or corrupted due to issues in data storage, transfer, or processing.
-
-Now, please provide the top 3 most likely reasons, order by likelihood (use your common sense to judge), in the following format:
-```json
-[{{"class": "The above 5 classes, or Other",
- "explanation": "Short explanation of the reason in 5 words",}}...]
-```"""
-
-
- messages = [{"role": "user", "content": template}]
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- messages.append(response['choices'][0]['message'])
- processed_string = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = json.loads(processed_string)
- self.document["missing_value"][column]["summary"] = json_code
- if LOG_MESSAGE_HISTORY:
- self.document["missing_value"][column]["history"] = messages
-
-
- def recommend_testing(self, overwrite=False, once=False):
- next_step = self.complete
-
- if "recommend_testing" in self.document:
- if self.document["recommend_testing"] and not overwrite:
- write_log("Warning: recommend_testing already exists in the document.")
- if not once:
- next_step()
- return
-
- create_progress_bar_with_numbers(3, doc_steps)
- print("🔍 Recommending testing...")
- progress = self.show_progress(1)
-
- basic_description = self.get_basic_description()
- table_name = self.get_table_name()
-
- json_code = recommend_testing(basic_description, table_name)
-
- progress.value += 1
-
- self.document["recommend_testing"] = json_code
-
- clear_output(wait=True)
-
- next_step()
-
- def check_unusual_all(self):
- next_step = self.recommend_testing
-
- column_meanings = self.document['column_meaning']['summary']
- unusual_columns = {}
- create_progress_bar_with_numbers(3, doc_steps)
- print("🤔 Checking the unusual values...")
- progress = self.show_progress(len(self.df.columns))
-
- for column_meaning in column_meanings:
- col = column_meaning['column']
- if col in self.name_mapping:
- col = self.name_mapping[col]
- meaning = column_meaning['meaning']
-
- self.check_unusual(col, meaning)
-
- result = self.document['unusual'][col]['summary']
-
- if result['Unusualness'] and 'Explanation' not in result:
- unusual_columns[col] = result
-
- progress.value += 1
-
- if not unusual_columns:
- next_step()
- return
-
- clear_output(wait=True)
- print("The following columns have unusual values. Please provide the explanation as much as possible.")
-
- def radio_change_handler(change):
- instance = change['owner']
- if instance.value == 'Explanation:':
- instance.text_area.layout.display = ''
- else:
- instance.text_area.layout.display = 'none'
-
- container = widgets.VBox()
-
-
- col_to_radio = {}
-
- for item in unusual_columns:
-
- reasons = unusual_columns[item]
-
- label_text = f"{item}: {reasons['Examples']}"
-
- label = widgets.HTML(value=label_text)
-
- options = ['Unclear', 'Explanation:']
-
- radio = widgets.RadioButtons(
- options=options,
- value=options[0],
- layout=widgets.Layout(width='80%', align_items='flex-start')
- )
-
- text_area = widgets.Textarea(
- value='',
- placeholder='Please provide the reason for unusual values.',
- description='',
- disabled=False,
- layout=widgets.Layout(display='none', width='100%')
- )
-
- radio.text_area = text_area
- col_to_radio[item] = radio
- radio.observe(radio_change_handler, names='value')
-
- item_container = widgets.VBox([label, radio, text_area])
- container.children += (item_container,)
-
- def submit_callback(btn):
- error_items = []
-
- for col in col_to_radio:
- radio = col_to_radio[col]
- if radio.value == 'Explanation:' and not radio.text_area.value.strip():
- error_items.append(col)
-
- if error_items:
- for col in error_items:
- print(f"\033[91m{col} explanation can't be empty.\033[0m")
- else:
- clear_output(wait=True)
-
- for col in col_to_radio:
- radio = col_to_radio[col]
-
- if radio.value == 'Explanation:':
- self.document["unusual"][col]["summary"]["Explanation"] = radio.text_area.value
- if LOG_MESSAGE_HISTORY:
- self.document["unusual"][col]["history"].append({"role": "user",
- "content": radio.text_area.value})
- else:
- self.document["unusual"][col]["summary"]["Explanation"] = "Unclear"
-
- print("Submission received.")
- next_step()
-
-
-
- submit_btn = widgets.Button(
- description="Submit",
- button_style='',
- tooltip='Submit',
- icon=''
- )
-
- submit_btn.on_click(submit_callback)
-
- display(container, submit_btn)
-
- if self.viewer:
- submit_callback(submit_btn)
-
-
-
- def check_unusual(self, col, meaning):
- if col not in self.df.columns:
- raise ValueError(f"Column {col} does not exist in the DataFrame.")
-
- if "unusual" not in self.document:
- self.document["unusual"] = {}
-
- if col not in self.document["unusual"]:
- self.document["unusual"][col] = {}
- else:
- if self.document["unusual"][col]:
- write_log(f"Warning: {col} unusual already exists in the document.")
- return
-
- df_col = self.df[col]
-
- if isinstance(df_col, pd.DataFrame):
- df_col = df_col.iloc[:, 0]
-
- unique_values = df_col.dropna().unique()
-
- values_string = f"The actual data have {len(unique_values)} unique values: "
-
- def construct_string_with_limit(base_string, values, char_limit):
- final_string = base_string
-
- included_values = []
-
- at_least_one = False
-
- for value in values:
- str_value = f"'{str(value)}'"
-
- if not at_least_one or len(final_string) + len(str_value) + len(included_values) * len(", ") < char_limit:
- included_values.append(str_value)
- at_least_one = True
- else:
- break
-
- values_part = ", ".join(included_values)
-
- if len(included_values) < len(values):
- values_part += "..."
-
- final_string += values_part
-
- return final_string
-
- char_limit = 300
- values_string = construct_string_with_limit(values_string, unique_values, char_limit)
-
-
-
- max_trials = 3
-
- while max_trials > 0:
-
- try:
-
- today = datetime.date.today()
-
- messages = [{"role": "user",
- "content": f"The column '{col}' is about: {meaning}. Guess in 10 words how the values usually look like."}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- template = f"""{values_string}
-
- Review if there are any unusual values. Look out for:
- 1. Values too large/small that are inconsistent with the context.
- E.g., age 999 or -5.
- Outlier is fine as long as it falls in a reasonable range, e.g., person age 120 is fine.
- 2. Patterns that don't align with the nature of the data.
- E.g., age 10.22
- 3. Special characters that don't fit within the real-world domain.
- E.g., age X21b
-
- Be careful about date as your knowledge of date is not updated. Today is {today}.
-
- Follow below step by step:
- 1. Summarize the values. Reason if it is unusual or also acceptable.
- 2. Conclude with the following dict:
-
- ```json
- {{
- "Unusualness": true/false,
- "Examples": "xxx values are unusual because ..." (empty if not unusual)
- }}
- ```"""
-
- messages.append(response['choices'][0]['message'])
- messages.append({"role": "user",
- "content": template})
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- messages.append(response['choices'][0]['message'])
-
- for message in messages:
- write_log(message['content'])
- write_log("-----------------------------------")
-
- processed_string = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = json.loads(processed_string)
-
- if "Unusualness" not in json_code:
- raise Exception("Unusualness is not in the json_code")
-
- if json_code["Unusualness"] and "Examples" not in json_code:
- raise Exception("Unusualness is true but Examples is not in the json_code")
-
- break
-
- except Exception as e:
- write_log(f"Error: {e}")
- write_log("-----------------------------------")
- max_trials -= 1
- if max_trials == 0:
- raise e
- continue
-
-
- self.document["unusual"][col]["summary"] = json_code
- if LOG_MESSAGE_HISTORY:
- self.document["unusual"][col]["history"] = messages
-
- def check_consistency_all(self):
- for col in self.df.columns:
- self.check_consistency(col)
-
- def check_consistency(self, col: str):
- if col not in self.df.columns:
- raise ValueError(f"Column {col} does not exist in the DataFrame.")
-
- if "consistency" not in self.document:
- self.document["consistency"] = {}
-
- if col not in self.document["consistency"]:
- self.document["consistency"][col] = {}
- else:
- if self.document["consistency"][col]:
- write_log(f"Warning: {col} already exists in the document.")
- return
-
- result = process_dataframe(self.df, 'consistency.txt', col_name=col)
-
- processed_string = escape_json_string(result[-1]["content"])
- content = json.loads(processed_string)['Action']['Content']
-
- if isinstance(content, str):
- self.document["consistency"][col]["summary"] = json.loads(content)
- else:
- self.document["consistency"][col]["summary"] = content
- if LOG_MESSAGE_HISTORY:
- self.document["consistency"][col]["history"] = result
-
- def check_pattern_all(self):
- for col in self.df.columns:
- self.check_pattern(col)
-
- def check_pattern(self, col: str):
- if col not in self.df.columns:
- raise ValueError(f"Column {col} does not exist in the DataFrame.")
-
- if pd.api.types.is_numeric_dtype(self.df[col].dtype):
- write_log(f"Warning: {col} is a numeric column. Skipping...")
- return
-
-
- if "pattern" not in self.document:
- self.document["pattern"] = {}
-
- if col not in self.document["pattern"]:
- self.document["pattern"][col] = {}
- else:
- if self.document["pattern"][col]:
- write_log(f"Warning: {col} already exists in the document.")
- return
-
- result = process_dataframe(self.df, 'pattern.txt', col_name=col)
-
- processed_string = escape_json_string(result[-1]["content"])
- content = json.loads(processed_string)['Action']['Content']
-
- if isinstance(content, str):
- self.document["pattern"][col]["summary"] = json.loads(content)
- else:
- self.document["pattern"][col]["summary"] = content
- if LOG_MESSAGE_HISTORY:
- self.document["pattern"][col]["history"] = result
-
- def to_yml(self):
- table_name = self.table_name
- description = self.document["main_entity"]["reasoning"]
- column_meanings = self.document["column_meaning"]["summary"]
-
- yaml_dict = {
- "version": 2,
- "models": [{
- "name": table_name,
- "description": description,
- "columns": []
- }]
- }
-
- for column_info in column_meanings:
- column_dict = {
- "name": column_info["column"],
- "description": column_info["meaning"]
- }
- yaml_dict["models"][0]["columns"].append(column_dict)
-
- return yaml_dict
-
- def write_yml_to_disk(self, filepath: str):
- yaml_dict = self.to_yml()
- with open(filepath, 'w') as file:
- yaml.dump(yaml_dict, file)
-
- def write_document_to_disk(self, filepath: str):
- with open(filepath, 'w') as file:
- json.dump(self.document, file)
-
- def read_document_from_disk(self, filepath: str, viewer=True):
- with open(filepath, 'r') as file:
- self.document = json.load(file)
- self.start_document(viewer=viewer)
-
- def __repr__(self):
- self.display_document()
- return ""
-
- def save_file(self):
- if self.table_name is not None:
- data_name = self.table_name
- else:
- data_name = self.document['main_entity']['summary']
-
- file_name = f"{data_name}_cocoon_data.json".replace(" ", "_")
-
- print(f"🤓 Do you want to save the file?")
-
- def save_file_click(b):
- updated_file_name = file_name_input.value
- allow_overwrite = overwrite_checkbox.value
-
- if os.path.exists(updated_file_name) and not allow_overwrite:
- print("\x1b[31m" + "Warning: Failed to save. File already exists." + "\x1b[0m")
- else:
- self.write_document_to_disk(updated_file_name)
- print(f"🎉 File saved successfully as {updated_file_name}")
-
- file_name_input = Text(value=file_name, description='File Name:')
-
- save_button = Button(description="Save File")
- save_button.on_click(save_file_click)
-
- overwrite_checkbox = Checkbox(value=False, description='Allow Overwrite')
-
- display(HBox([file_name_input, overwrite_checkbox]), save_button)
-
- if self.viewer:
- save_file_click(save_button)
-
-
-
-BOLD = '\033[1m'
-ITALIC = '\033[3m'
-END = '\033[0m'
-
-def get_value_from_path(data, path):
- """Recursively extract value from nested dict using the given path."""
-
- if not isinstance(data, dict):
- return data
-
- if not path:
- return data
- return get_value_from_path(data[path[0]], path[1:])
-
-target_concepts = {
- "PATIENT": {
- "Identity": {
- "Gender": ["gender_concept_id","gender_source_value"],
- "Race": ["race_concept_id","race_source_value"],
- "Ethnicity": ["ethnicity_concept_id","ethnicity_source_value"]
- },
- "Birth Details": {
- "Date of Birth": ["year_of_birth","month_of_birth","day_of_birth","birth_datetime"]
- },
- "Residence": {
- "Address": ["address","city","state","zip","county","location_source_value"],
- "Coordinates": ["latitude","longitude"]
- },
- "Identifier": ["person_id"]
- },
- "CONDITION_OCCURRENCE": {
- "Condition Information": {
- "Basic Condition Details": ["condition_occurrence_id", "condition_source_description_value", "condition_source_concept_id"],
- "Condition Duration": ["condition_start_date", "condition_start_datetime", "condition_end_date", "condition_end_datetime"],
- "Condition Status": ["stop_reason", "condition_status_source_value"]
- },
- "Patient Information": ["person_id"],
- "Provider Information": ["provider_id"],
- "Visit Information": {
- "Visit Identification": ["visit_occurrence_id", "visit_detail_id"],
- "Visit Context": ["condition_type_source"]
- }
- }
-
-}
-
-table_description = {
- "PATIENT": "The Patients table focuses on demographic and general information about individual patients. Key attributes include **person_id**, a unique identifier for each patient, and **gender_concept_id**, which records the patient's gender. The **year_of_birth**, **month_of_birth**, and **day_of_birth** attributes collectively provide the patient's date of birth. Other critical attributes are **race_concept_id** and **ethnicity_concept_id**, capturing the patient's race and ethnicity respectively. The **location_id** links to the patient's geographical information, while **provider_id** and **care_site_id** associate the patient with healthcare providers and care sites. Additionally, the **person_source_value**, **gender_source_value**, **gender_source_concept_id**, **race_source_value**, **race_source_concept_id**, **ethnicity_source_value**, and **ethnicity_source_concept_id** provide source-specific details for each corresponding attribute.",
- "VISIT_OCCURRENCE": "The VISIT_OCCURRENCE table documents detailed information about patient visits to healthcare facilities. It includes attributes such as **visit_occurrence_id** as a unique identifier for each visit, **person_id** linking to the patient's record, and **visit_concept_id** that classifies the type of visit (e.g., outpatient, inpatient). The table also tracks the **visit_start_date** and **visit_end_date** to specify the duration of the visit. Attributes like **visit_type_concept_id** offer insights into the context of the data collection, while **provider_id** links to the healthcare provider involved. Additionally, **care_site_id** associates the visit with a specific location, and **visit_source_value** and **visit_source_concept_id** capture information from original data sources. The table can also include **admitting_source_concept_id** and **discharge_to_concept_id** to detail patient transitions into and out of the care setting.",
- "CONDITION_OCCURRENCE": "The CONDITION_OCCURRENCE table primarily focuses on recording patient conditions, typically diagnosed by a healthcare provider. It contains key attributes such as **person_id**, which uniquely identifies a patient, and **condition_concept_id**, referencing the specific condition diagnosed. The attribute **condition_start_date** indicates when the condition was first observed, while **condition_end_date** reflects when it was resolved or ceased. Additionally, **condition_type_concept_id** describes the context or source of the diagnosis, such as a hospital visit. The **stop_reason** attribute provides insight into why a condition was considered resolved or ended. Another crucial attribute, **provider_id**, identifies the healthcare provider responsible for the diagnosis, and **visit_occurrence_id** links the condition to a specific patient visit. The table also includes **condition_source_value**, a textual representation of the condition from the original source data, and **condition_source_concept_id** for mapping to a standardized concept. Lastly, **condition_status_concept_id** gives further details about the status of the condition, like whether it's an active diagnosis or a historical record.",
- "DRUG_EXPOSURE": "The DRUG_EXPOSURE table captures information about a patient's exposure to a drug. The primary attribute, **DRUG_EXPOSURE_ID**, uniquely identifies each drug exposure event. **PERSON_ID** links to the individual patient, while **DRUG_CONCEPT_ID** identifies the specific drug. **DRUG_EXPOSURE_START_DATE** and **DRUG_EXPOSURE_END_DATE** denote the period of drug exposure. Dosage and frequency are detailed through **DOSE_UNIT_CONCEPT_ID** and **QUANTITY**. The table also includes **ROUTE_CONCEPT_ID** to specify the drug administration route, and **PROVIDER_ID** to identify the healthcare provider. **VISIT_OCCURRENCE_ID** links the drug exposure to a specific patient visit, and **DRUG_TYPE_CONCEPT_ID** categorizes the context of drug exposure, like prescription or inpatient administration.",
- "PROCEDURE_OCCURRENCE": "The PROCEDURE_OCCURRENCE table records details of clinical procedures performed on patients. At its core, this table includes the **procedure_concept_id**, which identifies the specific type of procedure carried out. The **person_id** attribute links the procedure to a specific patient. The **procedure_date** and **procedure_datetime** attributes specify when the procedure was performed. The **procedure_type_concept_id** defines the context or source of data entry. Other important attributes include **quantity**, representing the number of times the procedure was performed, and **provider_id**, identifying the healthcare provider who performed the procedure. The **visit_occurrence_id** links the procedure to the patient visit during which it was performed, and **modifier_concept_id** provides additional details or modifications to the procedure. Finally, **procedure_source_value** and **procedure_source_concept_id** offer source-specific codes for the procedure, and **qualifier_concept_id** gives further context or qualifiers related to the procedure.",
- "DEVICE": "The DEVICE table focuses on medical device usage in patient care. Key attributes include **device_id**, a unique identifier for each device record, and **person_id**, linking the device to an individual. **device_concept_id** represents the specific type of device. **device_exposure_start_date** and **device_exposure_end_date** define the usage period. **device_type_concept_id** describes the context of the device usage. **provider_id** identifies the healthcare provider involved, while **visit_occurrence_id** connects the device usage to a specific visit. Additional attributes like **device_source_value** and **quantity** provide extra details about the device and its usage.",
- "MEASUREMENT": "The MEASUREMENT table captures quantitative data and observations about a patient's health status. It includes key attributes like **measurement_id**, a unique identifier for each measurement, and **person_id**, linking the measurement to a specific patient. The **measurement_concept_id** identifies what was measured, such as blood pressure or glucose level. The **measurement_date** and **measurement_datetime** record when the measurement was taken. The **measurement_type_concept_id** indicates the nature of the measurement, like lab result or vital sign. The **operator_concept_id** represents how the measurement was taken, for example, by a device or a clinician. The **value_as_number**, **value_as_concept_id**, and **unit_concept_id** provide the result and unit of the measurement. **range_low** and **range_high** offer reference ranges for the measurement value. **provider_id** and **visit_occurrence_id** link the measurement to the provider and the visit during which it was taken. Additional details such as **measurement_source_value**, **measurement_source_concept_id**, and **unit_source_value** capture the original source information. Lastly, **value_source_value** records the measurement result as it was originally represented.",
- "OBSERVATION": "The OBSERVATION captures patient observations or measurements that are not diagnoses, procedures, or drug prescriptions. This table includes **person_id** to identify the patient, **observation_id** as a unique identifier for the observation, and **observation_concept_id** to specify the type of observation. The **observation_date** and **observation_datetime** record when the observation was made. **observation_type_concept_id** classifies the observation source, such as from a survey or a lab result. **value_as_number**, **value_as_string**, **value_as_concept_id**, and **unit_concept_id** describe the observation result in various formats. **qualifier_concept_id** and **associated_provider_id** give additional context to the observation, while **visit_occurrence_id** and **visit_detail_id** link the observation to specific visits. **obs_event_field_concept_id** captures the field from the source data where the observation originated, and **observation_source_value**, **observation_source_concept_id**, **unit_source_value**, and **qualifier_source_value** provide source-specific information. **observation_event_id**, **obs_event_field_concept_id**, and **value_as_datetime** offer further details on the observation event.",
- "DEATH": "The DEATH table is specifically designed to capture information about patient deaths. It contains several key attributes that provide details about the circumstances and documentation of a patient's death. The primary attribute is **person_id**, which uniquely identifies the deceased individual within the database. Another crucial attribute is **death_date**, specifying the exact date of death. The **death_datetime** attribute can offer more precise timing if available. The **death_type_concept_id** is used to indicate the source or method of death determination, like a death certificate or autopsy report. The **cause_concept_id** and **cause_source_value** attributes are used to record the cause of death, either as a standard concept or as a raw value from the data source, respectively. Additionally, the **cause_source_concept_id** represents a standardized concept that corresponds to the source value for the cause of death.",
- "SPECIMEN": "The SPECIMEN table stores detailed information about biological specimens collected from patients for diagnostic, treatment, or research purposes. This table includes several attributes that provide comprehensive data about each specimen. Key attributes include **specimen_id**, which is a unique identifier for each specimen, and **person_id**, linking the specimen to a specific individual in the database. The **specimen_concept_id** and **specimen_type_concept_id** describe the type of specimen and the method of its collection, respectively. **specimen_date** and **specimen_datetime** capture the date and time of specimen collection. **quantity** reflects the amount of specimen collected, and **unit_concept_id** indicates the unit of measurement. The **anatomic_site_concept_id** specifies the body site from where the specimen was taken, and **disease_status_concept_id** provides information about the disease status associated with the specimen. **specimen_source_id** and **visit_occurrence_id** are used to link the specimen to other relevant data in the database, such as the visit during which the specimen was collected. The **specimen_source_value**, **anatomic_site_source_value**, and **disease_status_source_value** attributes store the original values from the source data for mapping purposes.",
- "COST": "The COST table stores financial information related to healthcare services provided to patients. This table captures various cost-related details associated with healthcare encounters, procedures, and medications. Key attributes include **cost_event_id** and **cost_domain_id**, which identify the event and its domain (like drug or procedure) associated with the cost. The **currency_concept_id** specifies the currency of the cost. There's a focus on the specific nature of costs with attributes like **paid_copay**, **paid_coinsurance**, **paid_toward_deductible**, and **paid_by_payer**, detailing different payment components. The **amount_allowed** and **revenue_code_concept_id** provide information on allowable amounts and revenue codes.",
- "LOCATION": "The LOCATION table stores information about physical locations relevant to healthcare data. This table typically includes attributes such as **location_id**, which serves as a unique identifier for each location. It also contains **address_1** and **address_2** for detailed street addresses, **city** and **state** for regional identification, and **zip** for postal codes. Furthermore, the table includes **county**, **location_source_value**, and **country** to provide a comprehensive geographic context.",
- "CARE_SITE": "The CARE_SITE table represents information related to healthcare facilities where patient care is provided. It includes the **care_site_id**, a unique identifier for each care site. The **care_site_name** provides the name of the care site, and the **place_of_service_concept_id** links to a standardized concept identifying the type of care site (e.g., hospital, clinic). The **location_id** associates the care site with a physical location, and the **care_site_source_value** captures the original source data for the care site as it appears in the source system. Additionally, the **place_of_service_source_value** is the source code used in the source data to identify the type of care site.",
- "PROVIDER": "The PROVIDER tablecaptures detailed information about healthcare providers. It primarily focuses on the background and characteristics of individual providers. Key attributes include **provider_id**, serving as a unique identifier for each provider. The table also contains **provider_name**, which records the name of the provider. To classify the type of provider, there's **provider_type**, while **specialty_concept_id** links to their medical specialty. The **care_site_id** associates providers with their place of practice. Additionally, attributes like **gender_concept_id**, **year_of_birth**, and **provider_source_value** offer demographic and source-specific details about the providers.",
- "PAYER_PLAN_PERIOD": "The PAYER_PLAN_PERIOD table stores information about the periods of time a person is covered by a particular payer or plan. It captures details about the insurance or payer coverage for an individual. The table includes attributes such as **PAYER_PLAN_PERIOD_ID**, which is a unique identifier for each record. **PERSON_ID** links the record to a specific individual. The coverage period is defined by **PAYER_PLAN_PERIOD_START_DATE** and **PAYER_PLAN_PERIOD_END_DATE**, indicating the start and end of the coverage period. **PAYER_SOURCE_VALUE** and **PLAN_SOURCE_VALUE** provide information about the payer and plan from the source data. There are also attributes for standard concepts: **PAYER_CONCEPT_ID** and **PLAN_CONCEPT_ID**, along with their associated source values and source concept IDs (**PAYER_SOURCE_CONCEPT_ID** and **PLAN_SOURCE_CONCEPT_ID**). Lastly, **FAMILY_SOURCE_VALUE** may contain additional information about the family or group plan, if applicable."
-}
-
-
-table_samples = {
- "PATIENT":
-""" person_id gender_concept_id year_of_birth month_of_birth day_of_birth birth_datetime race_concept_id ethnicity_concept_id location_id provider_id care_site_id person_source_value gender_source_value gender_source_concept_id race_source_value race_source_concept_id ethnicity_source_value ethnicity_source_concept_id
-0 1 8507 1985 5 20 1985-05-20 8527 38003564 1 1 1 12345 M 8507 White 8527 Not Hispanic or Latino 38003564
-1 2 8532 1990 8 10 1990-08-10 8516 38003563 2 2 2 67890 F 8532 Black or African American 8516 Hispanic or Latino 38003563""",
-
- "VISIT_OCCURRENCE":
-"""visit_occurrence_id person_id visit_concept_id visit_start_date visit_start_datetime visit_end_date visit_end_datetime visit_type_concept_id provider_id care_site_id visit_source_value visit_source_concept_id admitting_source_concept_id admitting_source_value discharge_to_concept_id discharge_to_source_value preceding_visit_occurrence_id
-0 123456 101 9201 2023-01-10 2023-01-10 08:00:00 2023-01-12 2023-01-12 15:00:00 2001 501 601 VS123 3001 4001 AS123 5001 DS123 111
-1 789012 102 9202 2023-02-15 2023-02-15 09:30:00 2023-02-16 2023-02-16 10:30:00 2002 502 602 VS456 3002 4002 AS456 5002 DS456 222""",
-
- "CONDITION_OCCURRENCE":
-"""condition_occurrence_id person_id condition_concept_id condition_start_date condition_start_datetime condition_end_date condition_end_datetime condition_type_concept_id stop_reason provider_id visit_occurrence_id visit_detail_id condition_source_value condition_source_concept_id condition_status_concept_id condition_status_source_value
-0 1 101 201826 2021-01-10 2021-01-10 08:30:00 2021-01-20 2021-01-20 17:00:00 32019 Resolved 150 200 300 Diabetes 401 501 Active
-1 2 102 31967 2021-02-15 2021-02-15 09:45:00 2021-02-25 2021-02-25 16:30:00 32020 Improved 151 201 301 Hypertension 402 502 Controlled""",
-
- "DRUG_EXPOSURE":
-"""drug_exposure_id person_id drug_concept_id drug_exposure_start_date drug_exposure_end_date verbatim_end_date drug_type_concept_id stop_reason refills quantity days_supply sig route_concept_id lot_number provider_id visit_occurrence_id drug_source_value drug_source_concept_id route_source_value dose_unit_source_value
-0 1 1001 456789 2023-01-01 2023-01-10 2023-01-10 38000177 Completed course 0 10 10 Take 1 tablet daily 0 LOT1001 12345 111 DrugA 0 Oral mg
-1 2 1002 987654 2023-02-15 2023-03-01 2023-03-01 38000177 Adverse reaction 1 30 15 Take 2 tablets twice daily 0 LOT1002 67890 222 DrugB 0 Oral mg""",
-
- "PROCEDURE_OCCURRENCE":
-""" procedure_occurrence_id person_id procedure_concept_id procedure_date \
-0 1 101 2100001 2023-01-15
-1 2 102 2100002 2023-02-20
- procedure_datetime procedure_type_concept_id modifier_concept_id \
-0 2023-01-15 10:00:00 38000275 0
-1 2023-02-20 14:30:00 38000275 0
-
- quantity provider_id visit_occurrence_id visit_detail_id \
-0 1 501 701 0
-1 2 502 702 0
-
- procedure_source_concept_id procedure_source_value qualifier_concept_id \
-0 2100001 PROC1 0
-1 2100002 PROC2 0
-
- qualifier_source_value procedure_cost
-0 None 200.0
-1 None 450.0 """,
- "DEVICE":
-"""device_id person_id device_exposure_start_date device_exposure_start_datetime device_exposure_end_date device_exposure_end_datetime device_concept_id device_type_concept_id unique_device_id quantity provider_id visit_occurrence_id device_source_value device_source_concept_id
-0 1001 2001 2023-01-15 2023-01-15 08:00:00 2023-01-20 2023-01-20 18:00:00 3001 4001 UD1001 1 5001 6001 DeviceA 7001
-1 1002 2002 2023-02-20 2023-02-20 09:30:00 2023-02-25 2023-02-25 16:30:00 3002 4002 UD1002 2 5002 6002 DeviceB 7002""",
-
- "MEASUREMENT": """measurement_id person_id measurement_concept_id measurement_date measurement_datetime measurement_type_concept_id operator_concept_id value_as_number value_as_concept_id unit_concept_id range_low range_high provider_id visit_occurrence_id measurement_source_value measurement_source_concept_id unit_source_value value_source_value
-1001 2001 3001 2023-01-15 2023-01-15 08:00:00 4001 5001 5.6 6001 7001 4.5 8.0 8001 9001 BP_SYS 10001 mmHg 120
-1002 2002 3002 2023-01-16 2023-01-16 09:30:00 4002 5002 7.8 6002 7002 6.0 9.5 8002 9002 BP_DIA 10002 mmHg 80""",
-
- "OBSERVATION":
-"""observation_id person_id observation_concept_id observation_date \
-0 1001 123 3001 2023-01-01
-1 1002 456 3002 2023-01-02
-
- observation_datetime observation_type_concept_id value_as_number \
-0 2023-01-01 08:00:00 2001 98.6
-1 2023-01-02 09:30:00 2002 99.5
-
- value_as_string value_as_concept_id qualifier_concept_id unit_concept_id \
-0 Normal 4001 5001 6001
-1 Elevated 4002 5002 6002
-
- provider_id visit_occurrence_id observation_source_value \
-0 7001 8001 Blood Pressure
-1 7002 8002 Heart Rate
-
- observation_source_concept_id unit_source_value qualifier_source_value
-0 9001 mmHg Resting
-1 9002 beats/min After Exercise """,
-
- "DEATH":
-""" person_id death_date death_datetime death_type_concept_id cause_concept_id cause_source_value cause_source_concept_id
- 123456 2023-05-15 2023-05-15 14:30:00 38003565 50115 I21.9 4323456
- 789012 2023-06-20 2023-06-20 08:45:00 38003566 433146 C50.9 7654321""",
-
- "SPECIMEN":
-"""specimen_id person_id specimen_concept_id specimen_type_concept_id specimen_date specimen_datetime quantity unit_concept_id anatomic_site_concept_id disease_status_concept_id specimen_source_id visit_occurrence_id visit_detail_id specimen_source_value unit_source_value anatomic_site_source_value disease_status_source_value
-101 2001 3001 4001 2023-01-15 2023-01-15 10:00:00 1.5 5001 6001 7001 A123 8001 9001 Blood Sample ml Arm Healthy
-102 2002 3002 4002 2023-01-20 2023-01-20 11:30:00 2.0 5002 6002 7002 B456 8002 9002 Tissue Sample g Liver Diseased""",
-
- "COST":
-""" cost_id person_id cost_event_id cost_domain_id currency_concept_id total_charge total_paid payer_plan_period_id amount_allowed paid_by_payer paid_by_patient paid_patient_copay paid_patient_coinsurance paid_patient_deductible paid_by_primary paid_ingredient_cost paid_dispensing_fee
-0 1001 2001 3001 Drug 840 500.0 450.0 4001 450.0 400.0 50.0 25.0 15.0 10.0 400.0 300.0 20.0
-1 1002 2002 3002 Procedure 840 1500.0 1400.0 4002 1400.0 1300.0 100.0 50.0 40.0 10.0 1300.0 1000.0 50.0""",
-
- "LOCATION":
-"""location_id address_1 address_2 city state zip county location_source_value latitude longitude
-0 1 123 Main St Suite 100 Springfield NY 12345 Hampden L123 40.7128 -74.0060
-1 2 456 Elm St Apt 202 Riverdale CA 67890 Archie L456 34.0522 -118.2437""",
-
- "CARE_SITE":
-""" care_site_id care_site_name place_of_service_concept_id location_id care_site_source_value place_of_service_source_value
-1 City Hospital 12345 101 CS001 Hospital
-2 Rural Clinic 67890 102 CS002 Clinic""",
-
- "PROVIDER":
-"""provider_id provider_name npi dea specialty_concept_id care_site_id year_of_birth gender_concept_id provider_source_value specialty_source_value specialty_source_concept_id gender_source_value gender_source_concept_id
-0 101 Dr. Jane Smith 1234567890 AB1234567 111 10 1970 8507 P101 Cardiology 1001 F 1003
-1 102 Dr. John Doe 0987654321 CD7654321 222 20 1980 8507 P102 Neurology 1002 M 1004""",
-
- "PAYER_PLAN_PERIOD":
-"""payer_plan_period_id person_id payer_plan_period_start_date payer_plan_period_end_date payer_concept_id plan_concept_id sponsor_concept_id family_plan_concept_id stop_reason payer_source_value payer_source_concept_id plan_source_value plan_source_concept_id sponsor_source_value sponsor_source_concept_id family_plan_source_value family_plan_source_concept_id
-1001 123 2023-01-01 2023-12-31 2001 3001 4001 5001 End of contract PayerA 6001 PlanA 7001 SponsorA 8001 FamilyPlanA 9001
-1002 456 2023-06-01 2023-12-31 2002 3002 4002 5002 Change of employment PayerB 6002 PlanB 7002 SponsorB 8002 FamilyPlanB 9002""",
-}
-
-attributes_description = {
- "PATIENT" :{
- "person_id": "the original id from the source data provided, otherwise it can be an autogenerated number. ",
- "year_of_birth": "as an integer",
- "month_of_birth": "as an integer",
- "day_of_birth": "as an integer",
- "birth_datetime": "If birth_datetime is not provided in the source, use the following logic to infer the date: If day_of_birth is null and month_of_birth is not null then use the first of the month in that year. If month_of_birth is null or if day_of_birth AND month_of_birth are both null and the person has records during their year of birth then use the date of the earliest record, otherwise use the 15th of June of that year. If time of birth is not given use midnight (00:00:0000).",
- "gender_concept_id": "{ 'FEMALE': 8532, 'MALE': 8507 }",
- "gender_source_value": "the original value in the source table",
- "person_source_value": "any identifier from the source data that identifies the person.",
- "race_concept_id": "{ 'American Indian or Alaska Native': 8657, 'Asian': 8515, 'Black': 8516, 'Native Hawaiian or Other Pacific Islander': 8557, 'White': 8527 }",
- "race_source_value": "the original value in the source table",
- "ethnicity_concept_id": "{ 'Hispanic or Latino': 38003564, 'Not Hispanic or Latino': 38003563 }",
- "ethnicity_source_value": "the original value in the source table",
- },
- "CONDITION_OCCURRENCE": {
- "condition_occurrence_id": "the original id column from the source data provided, otherwise it can be an autogenerated number. ",
- "person_id": "The PERSON_ID of the PERSON for whom the condition is recorded.",
- "condition_start_date": "the start date of the condition",
- "condition_start_datetime": "If a source does not specify datetime the convention is to set the time to midnight (00:00:0000)",
- "condition_end_date": "The end date of the condition",
- "condition_end_datetime": " the end date of the condition",
- "stop_reason": "The Stop Reason indicates why a Condition is no longer valid with respect to the purpose within the source data.",
- "visit_occurrence_id": "The visit during which the condition occurred.",
- "condition_source_description_value": "Source data representing the condition that occurred.",
- "condition_type_source": "the provenance of the Condition record, as in whether the condition was from an EHR system, insurance claim, registry, or other sources.",
- "condition_status_source_value": "the source data indicating when and how a diagnosis was given to a patients",
- },
- "DEATH" :{
- "death_date": "Date of death, use January 1st of the year if only year is known. E.g., 2020-01-01",
- "death_datetime": "Exact date and time of death, default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "cause_source_value": "Original cause of death from source data. E.g., 'Heart Attack'",
- "death_record_source": "Source of death record (e.g., death certificate, hospital record). E.g., 'Death Certificate'",
- },
- "DRUG_EXPOSURE" :{
- "drug_exposure_id": "Unique identifier for each drug exposure record. E.g., 101",
- "drug_exposure_start_date": "Start date of the drug exposure. E.g., 2022-03-15",
- "drug_exposure_start_datetime": "Exact start date and time of the drug exposure. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "drug_exposure_end_date": "End date of the drug exposure. E.g., 2022-04-14",
- "drug_exposure_end_datetime": "Exact end date and time of the drug exposure. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "verbatim_end_date": "Original end date as recorded in the source data. E.g., 2022-04-14",
- "stop_reason": "Reason the drug exposure was stopped. E.g., 'Completed treatment'",
- "refills": "Number of refills prescribed. E.g., 2",
- "quantity": "Quantity of the drug prescribed. E.g., 30 (tablets)",
- "days_supply": "Number of days the drug supply is expected to last. E.g., 30",
- "sig": "Directions for use as specified by the prescriber. E.g., 'Take 1 tablet daily'",
- "lot_number": "Lot number of the drug. E.g., 'LN123456'",
- "drug_source_value": "Original value of the drug as in the source data. E.g., 'Aspirin'",
- "route_source_value": "Original value of the route of administration in the source data. E.g., 'Oral'",
- "dose_unit_source_value": "Unit of the dose as in the source data. E.g., 'mg'",
- },
- "PROCEDURE_OCCURRENCE" :{
- "procedure_occurrence_id": "Unique identifier for the procedure record. E.g., 456789",
- "procedure_date": "Date when the procedure was performed. E.g., 2023-05-20",
- "procedure_datetime": "Exact date and time of the procedure. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "procedure_type_concept_id": "Concept ID indicating the type of procedure (e.g., surgical, diagnostic). E.g., 44818701 (for 'Surgical procedure')",
- "modifier_concept_id": "Concept ID for any modifiers related to the procedure. E.g., 2000000 ('Laparoscopic')",
- "quantity": "Number of times the procedure was performed. E.g., 1",
- "procedure_source_value": "Original value of the procedure as in the source data. E.g., 'Appendectomy'",
- "qualifier_source_value": "Additional details about the procedure from the source data. E.g., 'Laparoscopic'",
- },
- "DEVICE" :{
- "device_exposure_id": "Unique identifier for each device exposure record. E.g., 999888",
- "device_exposure_start_date": "Start date of the device exposure. E.g., 2023-01-15",
- "device_exposure_start_datetime": "Exact start date and time of the device exposure. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "device_exposure_end_date": "End date of the device exposure. E.g., 2023-06-15",
- "device_exposure_end_datetime": "Exact end date and time of the device exposure. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "quantity": "Number of devices used or exposed to. E.g., 1",
- "device_source_value": "Original value of the device as in the source data. E.g., 'Pacemaker'",
- },
- "MEASUREMENT" :{
- "measurement_id": "Unique identifier for the measurement record. E.g., 123456",
- "measurement_date": "Date when the measurement was taken. E.g., 2023-08-10",
- "measurement_datetime": "Exact date and time the measurement was taken. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "value_as_number": "Numerical value of the measurement. E.g., 120 (for systolic blood pressure)",
- "range_low": "Lower limit of the normal range for the measurement. E.g., 90 (for systolic blood pressure)",
- "range_high": "Upper limit of the normal range for the measurement. E.g., 120 (for systolic blood pressure)",
- "measurement_source_value": "Original value of the measurement as in the source data. E.g., 'Blood Pressure'",
- "unit_source_value": "Original unit of the measurement as in the source data. E.g., 'mmHg'",
- },
- "OBSERVATION" :{
- "observation_id": "Unique identifier for each observation record. E.g., 789123",
- "observation_date": "Date of the observation. E.g., 2023-02-15",
- "observation_datetime": "Exact date and time of the observation. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "value_as_number": "Numerical value of the observation if applicable. E.g., 10 (for 'Cigarettes per day')",
- "value_as_string": "Textual value of the observation if applicable. E.g., 'Non-smoker'",
- "observation_source_value": "Original value of the observation as in the source data. E.g., 'Smoking status'",
- "unit_source_value": "Original unit of the observation as in the source data. E.g., 'Cigarettes per day'",
- },
- "SPECIMEN" :{
- "specimen_id": "Unique identifier for each specimen record. E.g., 123456",
- "specimen_date": "Date when the specimen was collected. E.g., 2023-03-01",
- "specimen_datetime": "Exact date and time when the specimen was collected. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "quantity": "Quantity of the specimen collected. E.g., 10 (ml for blood)",
- "specimen_source_value": "Original value of the specimen as in the source data. E.g., 'Blood Sample'",
- "unit_source_value": "Original unit of the specimen quantity as in the source data. E.g., 'ml'",
- },
- "COST" :{
- "total_charge": "Total charge for the event. E.g., 200.00",
- "total_cost": "Total cost for the event. E.g., 150.00",
- "total_paid": "Total amount paid for the event. E.g., 150.00",
- "paid_by_payer": "Amount paid by the payer (e.g., insurance). E.g., 100.00",
- "paid_by_patient": "Amount paid by the patient. E.g., 50.00",
- "paid_patient_copay": "Patient's copay amount. E.g., 20.00",
- "paid_patient_coinsurance": "Patient's coinsurance amount. E.g., 10.00",
- "paid_patient_deductible": "Patient's deductible amount. E.g., 20.00",
- "paid_by_primary": "Amount paid by primary payer. E.g., 100.00",
- "paid_ingredient_cost": "Cost of the ingredient for a drug. E.g., 30.00",
- "paid_dispensing_fee": "Dispensing fee paid. E.g., 5.00",
- "amount_allowed": "Amount allowed by the payer. E.g., 150.00",
- "cost_source_value": "Original value of the cost as in the source data. E.g., 'Procedure cost'",
- },
- "LOCATION" :{
- "location_id": "Unique identifier for each location record. E.g., 123",
- "address_1": "First line of the address. E.g., '123 Main St'",
- "address_2": "Second line of the address (if applicable). E.g., 'Apt 4'",
- "city": "City of the location. E.g., 'New York'",
- "state": "State of the location. E.g., 'NY'",
- "zip": "Zip code of the location. E.g., '10001'",
- "county": "County of the location. E.g., 'New York County'",
- "location_source_value": "Original value of the location as in the source data. E.g., '123 Main St, Apt 4, New York, NY, 10001'",
- "latitude": "a float, must be between -90 and 90.",
- "longitude": "a float, must be between -180 and 180.",
- },
- "CARE_SITE" :{
- "care_site_id": "Unique identifier for each care site record. E.g., 456",
- "care_site_name": "Name of the care site. E.g., 'Main Street Medical Center'",
- "care_site_source_value": "Original value of the care site as in the source data. E.g., 'Main Street Medical Center, Outpatient'",
- "place_of_service_source_value": "Original place of service as in the source data. E.g., 'Outpatient Hospital'",
- },
- "PROVIDER" :{
- "provider_id": "Unique identifier for each provider record. E.g., 789",
- "provider_name": "Name of the provider. E.g., 'Dr. Jane Smith'",
- "NPI": "National Provider Identifier. E.g., '1234567890'",
- "DEA": "Drug Enforcement Administration registration number. E.g., 'AB1234567'",
- "year_of_birth": "Year of birth of the provider. E.g., 1970",
- "provider_source_value": "Original value of the provider as in the source data. E.g., 'Dr. Jane Smith'",
- "specialty_source_value": "Original specialty of the provider as in the source data. E.g., 'Cardiology'",
- "gender_source_value": "Original gender of the provider as in the source data. E.g., 'Female'",
- },
- "PAYER_PLAN_PERIOD" :{
- "payer_plan_period_id": "Unique identifier for each payer plan period record. E.g., 12345",
- "payer_plan_period_start_date": "Start date of the coverage period. E.g., 2023-01-01",
- "payer_plan_period_end_date": "End date of the coverage period. E.g., 2023-12-31",
- "stop_reason": "Reason for the end of the coverage period. E.g., 'Change of employment'",
- "payer_source_value": "Original value of the payer as in the source data. E.g., 'Medicare'",
- "plan_source_value": "Original value of the plan as in the source data. E.g., 'Medicare Part D'",
- "sponsor_source_value": "Original value of the sponsor as in the source data. E.g., 'Government'",
- },
- "VISIT_OCCURRENCE" :{
- "visit_occurrence_id": "Unique identifier for each visit record. E.g., 456789",
- "visit_start_date": "Start date of the visit. E.g., 2023-04-15",
- "visit_start_datetime": "Exact start date and time of the visit. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "visit_end_date": "End date of the visit. E.g., 2023-04-15",
- "visit_end_datetime": "Exact end date and time of the visit. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "visit_source_value": "Original value of the visit as in the source data. E.g., 'Outpatient Visit'",
- },
- "OBSERVATION_PERIOD" :{
- "observation_period_id": "Unique identifier for each observation period record. E.g., 112233",
- "observation_period_start_date": "Start date of the observation period. E.g., 2023-01-01",
- "observation_period_start_datetime": "Exact start date and time of the observation period. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- "observation_period_end_date": "End date of the observation period. E.g., 2023-12-31",
- "observation_period_end_datetime": "Exact end date and time of the observation period. Default to midnight if time unknown. E.g., 2020-01-01 00:00:00",
- }
-}
-
-cdm_tables = """1. PATIENT: Individual's healthcare identity hub, capturing demographics, contact information, and unique medical identifiers.
-2. OBSERVATION_PERIOD: Clinical data timeframe, marking the start and end dates of health monitoring or treatment phases.
-3. VISIT_OCCURRENCE: Healthcare interaction records, detailing visit dates, types (e.g., inpatient, outpatient), locations, and involved healthcare professionals.
-4. CONDITION_OCCURRENCE: Medical condition logs, encompassing diagnosis dates, condition types, severity indicators, and observed symptoms.
-5. DRUG_EXPOSURE: Medication intake documentation, including drug names, dosages, administration routes, prescribing dates, and treatment durations.
-6. PROCEDURE_OCCURRENCE: Patient treatment actions, specifying procedure types, execution dates, purposes, and attending practitioners.
-7. DEVICE: Medical device utilization specifics, identifying device types, application or implantation dates, functional purposes, and locations.
-8. MEASUREMENT: Clinical test outcomes, containing types of tests, numerical or categorical results, units, and test dates.
-9. OBSERVATION: Diverse health-related facts, sourced from examinations, questionnaires, or procedural outcomes, inclusive of non-clinical information. This is general, and please prefer to use MEASUREMENT or CONDITION_OCCURRENCE if possible.
-10. DEATH: End-of-life data, specifying causes, dates, and circumstances surrounding a patient's death.
-11. SPECIMEN: Biological sample records, denoting sample types, collection dates, handling processes, and preservation methods.
-12. COST: Healthcare financials, enumerating costs associated with medical events like procedures, medications, visits, and equipment.
-13. LOCATION: Geographical pinpointing, detailing physical addresses or coordinates of healthcare facilities or patient residences.
-14. CARE_SITE: Health service points, defining types of facilities (e.g., hospitals, clinics), specializations, and operational scopes.
-15. PROVIDER: Practitioner profiles, listing professionals' credentials, specialties, roles, and contact information.
-16. PAYER_PLAN_PERIOD: Insurance coverage chronology, indicating enrollment spans, plan types, provided benefits, and payer information.
-"""
-
-
-
-class CDMTransformation:
- def __init__(self, doc_df, log_file_path='data_log.txt'):
-
- if isinstance(doc_df, DocumentedData):
- self.doc_df = doc_df
- self.pipeline = doc_df.generate_pipeline()
- elif isinstance(doc_df, DataCleaning):
- self.doc_df = doc_df.doc_df
- self.pipeline = doc_df.generate_pipeline()
-
-
-
- self.document = {}
-
- self.log_file_path = log_file_path
- database_description = """The database primarily focuses on healthcare data, structured around several interconnected entities. The central entity is the **PATIENT** table, which contains details about individuals receiving medical care. Their healthcare journey is tracked through the **VISIT_OCCURRENCE** table, which records each visit to a healthcare facility. The **CONDITION_OCCURRENCE** table details any diagnosed conditions during these visits, while the **DRUG_EXPOSURE** table captures information on medications prescribed to the patients.
-Procedures performed are logged in the **PROCEDURE_OCCURRENCE** table, and any medical devices used are listed in the **DEVICE** table. The **MEASUREMENT** table records various clinical measurements taken, and the **OBSERVATION** table notes any other relevant clinical observations.
-In cases where a patient passes away, the **DEATH** table provides information on the mortality. The **SPECIMEN** table tracks biological samples collected for analysis, and the **COST** table details the financial aspects of the healthcare services.
-The **LOCATION**, **CARE_SITE**, and **PROVIDER** tables offer contextual data, respectively detailing the geographical locations, healthcare facilities, and medical professionals involved in patient care. Lastly, the **PAYER_PLAN_PERIOD** table provides information on the patients' insurance coverage details and durations."""
- self.database_description = database_description
-
- tables = [
- "PATIENT", "VISIT_OCCURRENCE", "CONDITION_OCCURRENCE", "DRUG_EXPOSURE",
- "PROCEDURE_OCCURRENCE", "DEVICE", "MEASUREMENT", "OBSERVATION",
- "DEATH", "SPECIMEN", "COST", "LOCATION", "CARE_SITE", "PROVIDER",
- "PAYER_PLAN_PERIOD"
- ]
- self.tables = tables
-
- edges = [
- ("PATIENT", "VISIT_OCCURRENCE"),
- ("VISIT_OCCURRENCE", "CONDITION_OCCURRENCE"),
- ("VISIT_OCCURRENCE", "DRUG_EXPOSURE"),
- ("VISIT_OCCURRENCE", "PROCEDURE_OCCURRENCE"),
- ("VISIT_OCCURRENCE", "DEVICE"),
- ("VISIT_OCCURRENCE", "MEASUREMENT"),
- ("VISIT_OCCURRENCE", "OBSERVATION"),
- ("PATIENT", "DEATH"),
- ("VISIT_OCCURRENCE", "SPECIMEN"),
- ("VISIT_OCCURRENCE", "COST"),
- ("VISIT_OCCURRENCE", "LOCATION"),
- ("VISIT_OCCURRENCE", "CARE_SITE"),
- ("VISIT_OCCURRENCE", "PROVIDER"),
- ("PATIENT", "PAYER_PLAN_PERIOD")
- ]
- self.edges = edges
-
-
- def display_database(self):
- database_description = replace_asterisks_with_tags(self.database_description)
- display(HTML("
OMOP CDM
" + database_description))
- visualize_graph(self.tables, self.edges)
-
- def write_log(self, message: str):
- self.log_file = open(self.log_file_path, 'a')
- self.log_file.write(message + '\n')
- self.log_file.close()
-
- def complete(self):
- print("Congratulation! The transformation is complete. 🚀")
-
- def write_document_to_disk(self, filepath: str):
- with open(filepath, 'w') as file:
- json.dump(self.document, file)
-
- def read_document_from_disk(self, filepath: str):
- with open(filepath, 'r') as file:
- self.document = json.load(file)
-
- def show_progress(self, max_value):
- progress = widgets.IntProgress(
- value=1,
- min=0,
- max=max_value+1,
- step=1,
- description='',
- bar_style='',
- orientation='horizontal'
- )
-
- display(progress)
- return progress
-
- def start(self):
- self.get_main_table()
-
- def get_main_table(self):
-
- next_step = self.decide_main_table
-
- if "main_table" not in self.document:
- self.document["main_table"] = {}
- else:
- if self.document["main_table"]:
- write_log("Warning: main_table already exists in the document.")
- next_step()
- return
-
- print("🤓 Identifying target tables to map to...")
-
- progress = self.show_progress(1)
- source_table_description = self.doc_df.document["table_summary"]["summary"]
-
- summary, messages = find_target_table(source_table_description)
- progress.value += 1
-
- for message in messages:
- write_log(message['content'])
- write_log("-----------------------------------")
-
- for table in summary:
- if table not in self.tables:
- raise ValueError(f"Table {table} does not exist in the CDM.")
-
- self.document["main_table"]["summary"] = summary
- if LOG_MESSAGE_HISTORY:
- self.document["main_table"]["history"] = messages
-
- next_step()
-
-
-
-
-
-
-
-
-
-
-
-
-
- def concept_mapping(self, target_table):
-
- next_step = self.write_codes2
-
- if "concept_mapping" not in self.document:
- self.document["concept_mapping"] = {}
-
- if target_table in self.document["concept_mapping"]:
- write_log(f"Warning: {target_table} already exists in the document for concept_mapping.")
- else:
- print("🤓 Identifying the concept mapping...")
-
- progress = self.show_progress(1)
- source_table_description = self.doc_df.document["table_summary"]["summary"]
- source_table_sample = self.doc_df.get_sample_text()
- target_table_description = table_description[target_table]
- target_table_sample = table_samples[target_table]
- transform_reason = self.document["main_table"]["summary"][target_table]
-
- summary, messages = get_concept_mapping(source_table_description, source_table_sample, target_table_description, target_table_sample, transform_reason)
- progress.value += 1
-
- for message in messages:
- write_log(message['content'])
- write_log("-----------------------------------")
-
-
- self.document["concept_mapping"][target_table] = {}
- self.document["concept_mapping"][target_table]["summary"] = summary
- if LOG_MESSAGE_HISTORY:
- self.document["concept_mapping"][target_table]["history"] = messages
-
- summary = self.document["concept_mapping"][target_table]["summary"]
-
- print(f"""💡 {BOLD}Plan to map attributes from source to target table:{END}""")
-
-
- def display_mapping(summary):
- for mapping in summary:
- source_attributes = mapping["source_columns"]
- target_attributes = mapping["target_columns"]
- reason = mapping["reason"]
- print(f"{BOLD}{', '.join(source_attributes)}{END}")
- print(f" 🤓 Can be mapped to {BOLD}{', '.join(target_attributes)}{END}")
- print(f" {ITALIC}{reason}{END}")
-
- has_warning = False
- for source_attribute in source_attributes:
- warnings = self.doc_df.get_column_warnings(source_attribute)
- if warnings:
- has_warning = True
- print(f"\n ⚠️ {BOLD}{source_attribute}{END} has Data Quality Issues:")
- for idx, warning in enumerate(warnings):
- warning_type = warning["type"]
- warning_explanation = warning["explanation"]
- print(f" {idx+1}. {BOLD}{warning_type}{END}: {ITALIC}{warning_explanation}{END}")
- if has_warning:
- print(f" ⚠️ It's recommended to first clean the data before transformation.")
- print()
-
- display_mapping(summary)
-
- submit_button = widgets.Button(
- description='Next',
- disabled=False,
- button_style='',
- tooltip='Click to submit',
- )
-
- def on_submit_button_clicked(b):
- clear_output(wait=True)
- next_step(target_table)
-
-
- submit_button.on_click(on_submit_button_clicked)
-
- display(submit_button)
-
- print(f"""\n⚠️ Some attributes are not supported:""")
- print(f""" 1. 💭 concept id: vocabulary standardization is under development""")
- print(f""" 2. 🔗 foreign key: table connection is under development""")
- print(f"""😊 Please send a feature request if you want them!""")
-
- def write_codes2(self, target_table):
-
- print("💻 Writing the codes...")
-
- next_step = self.complete
-
- concept_mapping = self.document["concept_mapping"][target_table]["summary"]
-
- progress = self.show_progress(len(concept_mapping))
-
- for mapping in concept_mapping:
-
- target_attributes = mapping["target_columns"]
-
- potential_attributes = attributes_description[target_table]
-
- target_attributes = [attr for attr in target_attributes if attr in potential_attributes]
-
- if not target_attributes:
- progress.value += 1
- continue
-
- source_attributes = mapping["source_columns"]
-
- key = str(target_attributes) + str(source_attributes)
-
- reason = mapping["reason"]
-
- progress.value += 1
-
- if "code_mapping" not in self.document:
- self.document["code_mapping"] = {}
- else:
- if key in self.document["code_mapping"]:
- write_log(f"Warning: code_mapping for {key} already exists in the document.")
- continue
-
- source_table_description = self.doc_df.get_basic_description(sample_cols=source_attributes, cols=source_attributes)
-
- codes, messages = write_code_and_debug(key=key,
- source_attributes=source_attributes,
- source_table_description=source_table_description,
- target_attributes=target_attributes,
- df=self.doc_df.df,
- target_table=target_table)
-
- for message in messages:
- write_log(message['content'])
- write_log("-----------------------------------")
-
- self.document["code_mapping"][key] = {}
- self.document["code_mapping"][key]["summary"] = codes
- if LOG_MESSAGE_HISTORY:
- self.document["code_mapping"][key]["history"] = messages
-
- final_node_idx = self.pipeline.find_final_node()
-
- transform_step_indices = []
- project_step_indices = []
-
- sample_df = self.pipeline.run_codes()[:4]
-
- for mapping in concept_mapping:
- target_attributes = mapping["target_columns"]
-
- potential_attributes = attributes_description[target_table]
-
- target_attributes = [attr for attr in target_attributes if attr in potential_attributes]
-
- if not target_attributes:
- continue
-
- source_attributes = mapping["source_columns"]
-
- key = str(target_attributes) + str(source_attributes)
-
-
- code = self.document["code_mapping"][key]["summary"]
-
- transform_step = TransformationStep(name = "Column Transformation",
- explanation=f"""Map from source table {BOLD}{str(source_attributes)}{END} to target table {BOLD}{str(target_attributes)}{END}""",
- codes=code,
- sample_df=sample_df)
-
-
- step_index = self.pipeline.add_new_step(transform_step)
- transform_step_indices.append(step_index)
-
- project_step = ProjectionStep(name = "Column Projection",
- cols=target_attributes)
-
- step_index = self.pipeline.add_new_step(project_step)
- project_step_indices.append(step_index)
-
-
-
-
- concat_step = ConcatenateHorizontalStep()
- concat_step_idx = self.pipeline.add_new_step(concat_step)
-
-
- for transform_step_idx in transform_step_indices:
- self.pipeline.add_edge_by_index(final_node_idx, transform_step_idx)
-
- for i in range(len(transform_step_indices)):
- self.pipeline.add_edge_by_index(transform_step_indices[i],
- project_step_indices[i])
-
- for project_step_idx in project_step_indices:
- self.pipeline.add_edge_by_index(project_step_idx, concat_step_idx)
-
- self.pipeline.display()
-
- def print_codes(self):
- self.pipeline.print_codes()
-
- def run_codes(self):
- return self.pipeline.run_codes()
-
- def decide_one_one_table(self, target_table):
-
- next_step = self.concept_mapping
-
- if "one_to_one" not in self.document:
- self.document["one_to_one"] = {}
-
- if target_table in self.document["one_to_one"]:
- write_log(f"Warning: {target_table} already exists in the document for one_to_one.")
- else:
- print("🤓 Identifying the row mapping...")
-
- progress = self.show_progress(1)
-
-
- source_table_description = self.doc_df.document["table_summary"]["summary"]
- source_table_sample = self.doc_df.get_sample_text()
- target_table_description = table_description[target_table]
- target_table_sample = table_samples[target_table]
- transform_reason = self.document["main_table"]["summary"][target_table]
-
- summary, messages = decide_one_one(source_table_description, source_table_sample, target_table_description, target_table_sample, transform_reason)
- progress.value += 1
-
- for message in messages:
- write_log(message['content'])
- write_log("-----------------------------------")
-
- if not isinstance(summary["1:1"], bool):
- raise ValueError(f"summary['1:1'] is not a boolean.")
-
- self.document["one_to_one"][target_table] = {}
- self.document["one_to_one"][target_table]["summary"] = summary
- if LOG_MESSAGE_HISTORY:
- self.document["one_to_one"][target_table]["history"] = messages
-
- summary = self.document["one_to_one"][target_table]["summary"]
-
- if summary["1:1"]:
- next_step(target_table)
- else:
- display(HTML("☹️ Source doesn't have 1-1 row mapping with Target: " + summary["reason"]))
- print("😊 M-N row mapping is under development. Please send a feature request!")
-
-
-
- def decide_main_table(self):
-
- next_step = self.decide_one_one_table
-
-
- json_code = self.document["main_table"]["summary"]
-
-
- source_table_description = self.doc_df.document["table_summary"]["summary"]
- source_table_description = replace_asterisks_with_tags(source_table_description)
-
- display(HTML("Source table " + source_table_description))
-
-
- if json_code:
- print("💡 Below are potential tables to transform to.")
-
- for key in json_code:
- print(f' {BOLD}{key}{END}: {json_code[key]}')
-
- print("🤓 Please choose one table to transform to.")
-
- radio_options = widgets.RadioButtons(
- options=list(json_code.keys()) + ['Manually Specify'],
- description='',
- disabled=False
- )
- else:
- print(f"🙁 It doesn't seem to be related to any table in common data model. \n🤓 Please manually specify the \033[1mmost\033[0m related table.")
- radio_options = widgets.RadioButtons(
- options=['Manually Specify'],
- description='',
- disabled=False
- )
-
- dropdown = widgets.Dropdown(
- options=self.tables,
- description='',
- disabled= (True if json_code else False),
- layout={'display': 'none' if json_code else ''}
- )
-
- def on_radio_selection_change(change):
- if change['new'] == 'Manually Specify':
- dropdown.disabled = False
- dropdown.layout.display = ''
- else:
- dropdown.disabled = True
- dropdown.layout.display = 'none'
-
- radio_options.observe(on_radio_selection_change, names='value')
-
- submit_button = widgets.Button(
- description='Submit',
- disabled=False,
- button_style='',
- tooltip='Click to submit',
- )
-
- def on_submit_button_clicked(b):
- if radio_options.value == 'Manually Specify':
- main_table = dropdown.value
- else:
- main_table = radio_options.value
-
- clear_output(wait=True)
-
- next_step(main_table)
-
-
- submit_button.on_click(on_submit_button_clicked)
-
- container = widgets.VBox([radio_options, dropdown, submit_button])
-
- display(container)
-
- def write_codes(self):
-
- next_step = self.complete
-
- print("💻 Writing the codes...")
-
- return
-
- concept_mapping = self.document["concept_mapping"]
-
- source_concepts = self.doc_df.document["column_grouping"]["summary"]
-
- target_to_source = {}
-
- for key, value in concept_mapping.items():
- source_path = key.strip("[]").replace("'", "").split(", ")
-
- if "summary" in value:
- for summary_mapping in value["summary"]:
- target_path = str(summary_mapping)
- if target_path not in target_to_source:
- target_to_source[target_path] = source_path
- else:
- target_to_source[target_path].append(source_path)
-
- progress = self.show_progress(len(target_to_source))
-
- for key, source_path in target_to_source.items():
- target_path = key.strip("[]").replace("'", "").split(", ")
- target_attributes = get_value_from_path(target_concepts, target_path)
- source_attributes = get_value_from_path(source_concepts, source_path)
-
- progress.value += 1
-
- if "code_mapping" not in self.document:
- self.document["code_mapping"] = {}
- else:
- if key in self.document["code_mapping"]:
- write_log(f"Warning: code_mapping for {key} already exists in the document.")
- continue
-
- self.write_code_single(key, source_attributes, target_attributes)
-
- self.target_df = pd.DataFrame()
-
- for key, source_path in target_to_source.items():
- target_path = key.strip("[]").replace("'", "").split(", ")
- target_attributes = get_value_from_path(target_concepts, target_path)
- source_attributes = get_value_from_path(source_concepts, source_path)
-
-
- print(f"""Codes that map
-# from source table {BOLD}{str(source_path)}{END} ({ITALIC}{str(source_attributes)}{END})
-# to target table {BOLD}{str(target_path)}{END} ({ITALIC}{str(target_attributes)}{END})""")
-
- code = self.document["code_mapping"][key]
- print()
- print("-" * 80)
- print(highlight(code, PythonLexer(), Terminal256Formatter()))
- print("-" * 80)
- print()
-
-
- exec(code, globals())
- temp_target_df = etl(self.doc_df.df)
- for col in temp_target_df.columns:
- self.target_df[col] = temp_target_df[col]
-
-
-
- def write_code_single(self, key, source_attributes, target_attributes):
-
- target_attributes = "\n".join(f"{idx + 1}. {target_attribute}: {attributes_description[target_attribute]}" for idx, target_attribute in enumerate(target_attributes))
-
- template = f"""ETL task: Given Source Table, tansform it into Target Table with new columns.
-
-Source table:
-{self.doc_df.get_basic_description(sample_cols=source_attributes, cols=source_attributes)}
-
-The target table needs columns:
-{target_attributes}
-
-Do the following:
-1. First reason about how to extract the columns
-2. Then fill in the python function with detailed comments.
-```python
-def etl(source_df):
- target_df = pd.DataFrame()
- ...
- return target_df
-```"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- python_code = extract_python_code(response['choices'][0]['message']['content'])
-
- detailed_error_info = None
-
- max_tries = 2
-
- while max_tries > 0:
- max_tries -= 1
-
- try:
- exec(python_code, globals())
- temp_target_df = etl(self.doc_df.df)
- except Exception:
- detailed_error_info = get_detailed_error_info()
-
- if detailed_error_info is None:
- self.document["code_mapping"][key] = python_code
- return
-
-
- error_message = f"""There is a bug in the code: {detailed_error_info}.
-First, study the error message and point out the problem.
-Then, fix the bug and return the codes in the following format:
-```python
-def etl(source_df):
- target_df = pd.DataFrame()
- ...
- return target_df
-```"""
- messages = [{"role": "user", "content":template},
- {"role": "assistant", "content": python_code},
- {"role": "user", "content": error_message},]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(error_message)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- python_code = extract_python_code(response['choices'][0]['message']['content'])
-
- raise Exception("The code is not correct. Please try again.")
-
-
- def map_concept(self):
-
- next_step = self.write_codes
-
- print("Mapping the concepts...")
-
- target_concept = {self.document["main_table"]["summary"]["concept"]:
- target_concepts[self.document["main_table"]["summary"]["concept"]]}
-
- def dict_to_descriptive_list(data_dict, parent_keys=None):
- if parent_keys is None:
- parent_keys = []
-
- descriptive_map = {}
-
- for key, value in data_dict.items():
- current_keys = parent_keys + [key]
-
- if isinstance(value, dict):
- descriptive_map.update(dict_to_descriptive_list(value, current_keys))
- elif isinstance(value, list):
- keys_path = '[' + ', '.join(current_keys) + ']'
-
- attributes = ', '.join(value)
-
- description = f"{keys_path}, with {len(value)} attributes: {attributes}"
-
- descriptive_map[str(current_keys)] = description
-
- return descriptive_map
-
- source_concept = dict_to_descriptive_list(self.doc_df.document["column_grouping"]["summary"])
-
- progress = self.show_progress(len(source_concept))
-
- for keys_path in source_concept:
- description = source_concept[keys_path]
-
- if "concept_mapping" not in self.document:
- self.document["concept_mapping"] = {}
- else:
- if keys_path in self.document["concept_mapping"]:
- write_log(f"Warning: concept_mapping for {keys_path} already exists in the document.")
- continue
-
- self.map_concept_single(keys_path, description, target_concept)
-
- progress.value += 1
-
- clear_output(wait=True)
-
- source_concepts = self.doc_df.document["column_grouping"]["summary"]
-
- concept_mapping = self.document["concept_mapping"]
-
-
-
- def display_mapping(concept_mapping, target_concepts, source_concepts):
- results = []
-
- for key, value in concept_mapping.items():
- source_path = key.strip("[]").replace("'", "").split(", ")
-
- source_attributes = get_value_from_path(source_concepts, source_path)
-
- print(f"{BOLD}{'->'.join(source_path)}{END} ({ITALIC}{', '.join(source_attributes)}{END})")
- if "summary" in value:
- for summary_mapping in value["summary"]:
- target_path = summary_mapping
- target_attributes = get_value_from_path(target_concepts, target_path)
- matching_source_attributes = ", ".join(summary_mapping[-3:])
- print(f" Can be mapped to {BOLD}{'->'.join(target_path)}{END} ({ITALIC}{', '.join(target_attributes)}{END}) attributes.")
- else:
- print(f" Can't be used for any attributes.")
-
-
- display_mapping(concept_mapping, target_concepts, source_concepts)
-
- submit_button = widgets.Button(
- description='Next',
- disabled=False,
- button_style='',
- tooltip='Click to submit',
- )
-
- def on_submit_button_clicked(b):
-
- clear_output(wait=True)
-
- next_step()
-
-
- submit_button.on_click(on_submit_button_clicked)
-
- display(submit_button)
-
- def map_concept_single(self, keys_path, description, target_concept):
-
- key_path_list = keys_path.strip("[]").replace("'", "").split(", ")
-
- def extract_paths(data, path=None, results=None):
- if path is None:
- path = []
- if results is None:
- results = []
-
- for key, value in data.items():
- new_path = path + [key]
- if isinstance(value, dict):
- extract_paths(value, new_path, results)
- elif isinstance(value, list):
- results.append(new_path)
-
- return results
-
- paths = extract_paths(target_concept)
-
- paths_str = "\n".join(f"{idx + 1}. {item}" for idx, item in enumerate(paths))
-
- template = f"""You have list of target concepts about {list(target_concept.keys())[0]}. Each concept is a list from category to specifics:
-{paths_str}
-
-You have a source table about: {description}.
-The goal is to transform from source to target.
-
-Enumerate the target concepts that the source table can be potentially mapped to, in the following format (empty list if no relevant):
-```json
-[["{list(target_concept.keys())[0]}",...,"Leaf Category"],
- ["{list(target_concept.keys())[0]}", ...]...]
-```"""
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- processed_string = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = json.loads(processed_string)
-
-
-
- if len(json_code) == 0:
- self.document["concept_mapping"][keys_path] = {}
- return
-
-
-
- result = "\n".join(f'{idx + 1}. {path}: {get_value_from_path(target_concept, path)}' for idx, path in enumerate(json_code))
-
- template = f"""You have list of target concepts about {list(target_concept.keys())[0]}. Each concept is a list from category to specifics:
-{result}
-
-You have a source table about: {description}.
-**Assumption: {list(target_concept.keys())[0]} is semantically similar or more general to {key_path_list[0]}!!**
-
-Exclude target concepts that are obviously semantically different.
-E.g., ["Person", "birth date"] and ["Patient, "death date"]
-"birth date" and "death date", despite both are about date, are obviously different.
-Sometimes, the difference is not obvious. E.g., "birth date" and "birth year" are the same.
-
-Based on the assumption, exclude the conpets that are obviously different, but keep the ones that are unsure.
-Return the remaining concepts in the following format (empty list if no).
-```json
-[["{list(target_concept.keys())[0]}",...,"Leaf Category"],
- ["{list(target_concept.keys())[0]}", ...]...]
-```"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- processed_string = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = json.loads(processed_string)
-
-
-
- if len(json_code) == 0:
- self.document["concept_mapping"][keys_path] = {}
- return
-
-
-
-
- result = "\n".join(f'{idx + 1}. {path}: {get_value_from_path(target_concept, path)}' for idx, path in enumerate(json_code))
-
- template = f"""You have list of target concepts about {list(target_concept.keys())[0]}. Each concept is a list from category to specifics. The right side is its attributes:
-{result}
-
-You have a source table about: {description}.
-First, go through the attributes of target concept. Argue if any can be transformed from the source table.
-E.g., "Student height" cannot be transformed from "Student weight" as they are different measurements.
-E.g., "Student age" can be transformed from "Student birth date" by calculating the date difference.
-
-Then, enumerate the target concept where there exists attributes can be transformed (empty list if no):
-```json
-[["{list(target_concept.keys())[0]}",...,"Leaf Category"],
- ["{list(target_concept.keys())[0]}", ...]...]
-```"""
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- processed_string = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = json.loads(processed_string)
-
-
-
- self.document["concept_mapping"][keys_path] = {}
-
- if len(json_code) == 0:
- return
-
- self.document["concept_mapping"][keys_path]["summary"] = json_code
-
-
-
-
-
-def check_functional_dependency(df, determinant, dependent):
- groups = df.groupby(list(determinant))[dependent].nunique()
- is_functionally_dependent = (groups == 1).all()
- return is_functionally_dependent
-
-
-
-def combine_pipelines(pipelines):
- new_steps = []
- new_edges = {}
- offset = 0
-
- for pipeline in pipelines:
- new_steps.extend(pipeline.steps)
-
- for source_idx, targets in pipeline.edges.items():
- new_targets = [t + offset for t in targets]
- new_edges[source_idx + offset] = new_targets
-
- offset += len(pipeline.steps)
-
- return TransformationPipeline(new_steps, new_edges)
-
-def find_instance_index(instance_list, target_instance):
- for idx, instance in enumerate(instance_list):
- if instance is target_instance:
- return idx
- return -1
-
-
-
-
-def find_final_node(steps, edges):
-
- node_set = set(range(len(steps)))
-
- nodes_with_no_outgoing = {node for node, targets in edges.items() if len(targets) == 0}
-
- nodes_not_in_edges = node_set - set(edges.keys())
-
- potential_final_nodes = nodes_with_no_outgoing.union(nodes_not_in_edges)
-
- if len(potential_final_nodes) != 1:
- return None
-
- final_node = potential_final_nodes.pop()
-
- for node in node_set:
- if node == final_node:
- continue
- if not find_path(edges, node, final_node):
- return None
-
- return final_node
-
-
-def find_source_node(nodes, edges):
- if not edges and len(nodes) == 1:
- return 0
-
- all_nodes_with_outgoing = set(edges.keys())
-
- nodes_with_incoming = set()
- for targets in edges.values():
- nodes_with_incoming.update(targets)
-
- source_candidates = all_nodes_with_outgoing - nodes_with_incoming
-
- if len(source_candidates) == 1:
- return source_candidates.pop()
- else:
- return None
-
-
-
-
-def find_path(edges, start, end, visited=None):
- if visited is None:
- visited = set()
-
- if start == end:
- return True
- if start in visited:
- return False
-
- visited.add(start)
-
- for neighbor in edges.get(start, []):
- if find_path(edges, neighbor, end, visited):
- return True
-
- return False
-
-
-
-
-
-class TransformationPipeline:
-
-
- def __init__(self, steps, edges):
- if not isinstance(steps, list):
- raise ValueError("Steps must be a list")
-
- for step in steps:
- if not isinstance(step, TransformationStep):
- raise ValueError("Each step must be an instance of TransformationStep")
-
- if not isinstance(edges, dict):
-
- if not isinstance(edges, list):
- raise ValueError("Edges must be a list")
-
- for edge in edges:
- if not isinstance(edge, tuple) or len(edge) != 2:
- raise ValueError("Each edge must be a tuple of two elements")
- if not isinstance(edge[0], int) or not isinstance(edge[1], int):
- raise ValueError("Edge elements must be integers")
- if edge[0] < 0 or edge[0] >= len(steps) or edge[1] < 0 or edge[1] >= len(steps):
- raise ValueError("Edge indices must be within the range of steps")
-
- self.edges = {}
- for source_idx, target_idx in edges:
- if source_idx not in self.edges:
- self.edges[source_idx] = []
- self.edges[source_idx].append(target_idx)
-
- else:
- self.edges = edges
-
- self.steps = steps
-
- self.cached_results = {}
-
-
-
- def get_step(self, idx):
- return self.steps[idx]
-
- def get_step_idx(self, step):
- return find_instance_index(self.steps, step)
-
-
-
-
-
-
-
-
-
- def validate_graph(self):
- if self._is_cyclic():
- raise ValueError("The graph is cyclic.")
-
- num_steps = len(self.steps)
- for source, targets in self.edges.items():
- if source >= num_steps or any(target >= num_steps for target in targets):
- raise ValueError("Edge indices are out of range.")
-
- def _is_cyclic(self):
- visited = set()
- rec_stack = set()
-
- def _is_cyclic_util(v):
- if v not in visited:
- visited.add(v)
- rec_stack.add(v)
-
- for neighbour in self.edges.get(v, []):
- if neighbour not in visited and _is_cyclic_util(neighbour):
- return True
- elif neighbour in rec_stack:
- return True
-
- rec_stack.remove(v)
- return False
-
- for node in range(len(self.steps)):
- if node not in visited and _is_cyclic_util(node):
- return True
-
- return False
-
- def display_workflow(self):
- nodes = [f"{idx + 1}. {step.name}" for idx, step in enumerate(self.steps)]
-
- display_workflow(nodes, self.edges)
-
- def get_nodes(self):
- nodes = [f"{idx + 1}. {step.name}" for idx, step in enumerate(self.steps)]
- return nodes
-
- def display(self, call_back=None):
-
- if call_back is None:
- call_back = self.display
-
-
- def create_widget(instances):
-
- nodes = self.get_nodes()
-
- dropdown = widgets.Dropdown(
- options=nodes,
- disabled=False,
- )
-
- button1 = widgets.Button(description="View")
- button2 = widgets.Button(description="Edit")
-
- button3 = widgets.Button(description="Return")
-
- def on_button_clicked3(b):
- clear_output(wait=True)
- call_back()
-
- def on_button_clicked(b):
- clear_output(wait=True)
-
- display_workflow(nodes, self.edges)
-
- display(dropdown)
-
- buttons = widgets.HBox([button1, button2, button3])
- display(buttons)
-
- idx = nodes.index(dropdown.value)
-
- selected_instance = instances[idx]
- selected_instance.display()
-
- def on_button_clicked2(b):
- clear_output(wait=True)
-
- idx = nodes.index(dropdown.value)
- selected_instance = instances[idx]
-
- def call_back_display(step):
- clear_output(wait=True)
- call_back()
-
- selected_instance.edit_widget(callbackfunc=call_back_display)
-
- button1.on_click(on_button_clicked)
- button2.on_click(on_button_clicked2)
- button3.on_click(on_button_clicked3)
-
- display_workflow(nodes, self.edges)
-
- buttons = widgets.HBox([button1, button2])
-
- display(dropdown, buttons)
-
-
- create_widget(self.steps)
-
-
-
-
- def get_codes(self):
- sorted_step_idx = topological_sort(self.steps, self.edges)
-
- codes = "from cocoon_data import *\n\n"
-
- for step_idx in sorted_step_idx:
-
- step = self.steps[step_idx]
-
- source_ids = get_source_nodes_ids(self.edges, step_idx)
-
- codes += step.get_codes(target_id=step_idx, source_ids=source_ids)
-
- return codes
-
-
- def run_codes(self, use_cache=True):
-
- sorted_step_idx = topological_sort(self.steps, self.edges)
-
- for step_idx in sorted_step_idx:
-
- if use_cache and step_idx in self.cached_results:
- continue
-
- step = self.steps[step_idx]
-
- source_ids = get_source_nodes_ids(self.edges, step_idx)
-
- source_dfs = [self.cached_results[source_id] for source_id in source_ids]
-
- result = step.run_codes(dfs=source_dfs)
-
- if isinstance(result, str):
- write_log(result)
- raise ValueError(result)
- else:
- self.cached_results[step_idx] = result
-
- final_node_idx = self.find_final_node()
- if final_node_idx is None:
- raise ValueError("No final node to return")
- else:
- return self.cached_results[final_node_idx]
-
-
- def print_codes(self):
- codes = self.get_codes()
-
- print(highlight(codes, PythonLexer(), Terminal256Formatter()))
-
-
-
- def find_final_node(self):
- return find_final_node(self.steps, self.edges)
-
- def get_final_step(self):
- final_node = self.find_final_node()
- final_step = self.get_step(final_node)
- return final_step
-
- def find_source_node(self):
- return find_source_node(self.steps, self.edges)
-
- def get_source_step(self):
- source_node = self.find_source_node()
- source_step = self.get_step(source_node)
- return source_step
-
- def remove_final_node(self):
- if len(self.steps) <= 1:
- raise ValueError("The graph has less than or equal to 1 step. Cannot remove the final node.")
-
-
- final_node = self.find_final_node()
- if final_node is None:
- raise ValueError("No final node to remove")
-
- self.steps.pop(final_node)
-
- updated_edges = {}
- for source, targets in self.edges.items():
- adjusted_targets = [t - 1 if t > final_node else t for t in targets if t != final_node]
-
- adjusted_source = source - 1 if source > final_node else source
-
- if adjusted_targets:
- updated_edges[adjusted_source] = adjusted_targets
-
- self.edges = updated_edges
-
- new_cache = {}
- for key, value in self.cached_results.items():
- new_key = key - 1 if key > final_node else key
-
- if key != final_node:
- new_cache[new_key] = value
-
- self.cached_results = new_cache
-
- def add_step(self, new_step, parent_node_idx=None):
-
- if len(self.steps) == 0 and parent_node_idx is None:
- self.steps.append(new_step)
- return
-
- if parent_node_idx is not None and isinstance(parent_node_idx, TransformationStep):
- parent_node_idx = find_instance_index(self.steps, parent_node_idx)
-
- if parent_node_idx is not None and (parent_node_idx < 0 or parent_node_idx >= len(self.steps)):
- raise ValueError(f"Parent node index is out of range. {parent_node_idx} is not within [0, {len(self.steps) - 1}].")
-
- if not isinstance(new_step, TransformationStep):
- raise ValueError("The new step must be an instance of TransformationStep.")
-
- self.steps.append(new_step)
-
- new_step_idx = len(self.steps) - 1
-
- if parent_node_idx is not None:
- if parent_node_idx not in self.edges:
- self.edges[parent_node_idx] = []
- self.edges[parent_node_idx].append(new_step_idx)
-
-
-
- def add_step_right_after(self, new_step, parent_node_idx):
-
- if isinstance(parent_node_idx, TransformationStep):
- parent_node_idx = find_instance_index(self.steps, parent_node_idx)
-
- if parent_node_idx < 0 or parent_node_idx >= len(self.steps):
- raise ValueError(f"Parent node index is out of range. {parent_node_idx} is not within [0, {len(self.steps) - 1}].")
-
- if not isinstance(new_step, TransformationStep):
- raise ValueError("The new step must be an instance of TransformationStep.")
-
- self.steps.append(new_step)
-
- new_step_idx = len(self.steps) - 1
-
-
- children = self.edges.get(parent_node_idx, [])
- self.edges[parent_node_idx] = [new_step_idx]
- self.edges[new_step_idx] = children
-
- self.cached_results = {}
-
-
-
-
- def add_step_to_final(self, new_step):
-
- final_node = self.find_final_node()
- if final_node is None:
- raise ValueError("The graph doesn't have a final node.")
-
- self.add_step(new_step, final_node)
-
-
-
-
- def add_edge(self, source_step, target_step):
- source_idx = find_instance_index(self.steps, source_step)
- target_idx = find_instance_index(self.steps, target_step)
-
- if source_idx == -1:
- raise ValueError("The source step is not in the graph.")
- if target_idx == -1:
- raise ValueError("The target step is not in the graph.")
-
- self.add_edge_by_index(source_idx, target_idx)
-
-
-
- def add_new_step(self, new_step):
- self.steps.append(new_step)
- return len(self.steps) - 1
-
- def add_edge_by_index(self, source_idx, target_idx):
- if source_idx not in self.edges:
- self.edges[source_idx] = []
- self.edges[source_idx].append(target_idx)
-
- self.cached_results = {}
-
-
- def start(self):
- pass
-
- def __repr__(self):
- self.display()
- return ""
-
-class Testing:
- def __init__(self, name = None, explanation="", codes="", sample_df=None):
- if name is None:
- self.name = "Testing Task"
- else:
- self.name = name
-
- self.codes = codes
- self.explanation = explanation
- self.sample_df = sample_df
-
-
- def run_codes(self, dfs, codes = None):
-
- if codes is None:
- codes = self.codes
-
- if isinstance(dfs, list):
- df = dfs[0]
- else:
- df = dfs
-
- input_df = df.copy()
- try:
- if 'transform' in globals():
- del globals()['transform']
- exec(codes, globals())
- result = test(input_df)
- return result
-
- except Exception:
- detailed_error_info = get_detailed_error_info()
- return detailed_error_info
-
- def generate_codes(self, explanation=None):
- if explanation is None:
- explanation = self.explanation
- template = f"""Testing task: Given input df, write python codes that test df and output True/False
-===
-Input Df:
-{self.sample_df.to_csv()}
-===
-Testing Requirement
-{explanation}
-
-Do the following:
-1. First reason about how to test
-2. Then fill in the python function, with detailed comments.
-DONT change the function name and the return clause.
-
-{{
- "reason": "To transform, we need to ...",
- "codes": "def test(input_df):\\n ...\\n return True/False"
-}}
-"""
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
- self.reason = json_code["reason"]
- self.codes = json_code["codes"]
-
- def display(self):
- raise NotImplementedError
-
-
-class TestColumnUnique(Testing):
- def __init__(self, col_name, name = None, explanation="", codes="", sample_df=None):
- super().__init__(name, explanation, codes, sample_df)
- self.col_name = col_name
-
- if name is None:
- self.name = f"Test if {col_name} is unique"
-
- def generate_test_unique_codes(self):
- self.codes = """# Concatenate all dataframes horizontally
-def test(input_df):
- return input_df["{col_name}"].is_unique"""
-
-class TestColumnNotNull(Testing):
- def __init__(self, col_name, name=None, explanation="", codes="", sample_df=None):
- super().__init__(name, explanation, codes, sample_df)
- self.col_name = col_name
-
- if name is None:
- self.name = f"Test if {col_name} has no null"
-
- def generate_test_not_null_codes(self):
- self.codes = f"""def test(input_df):
- return not input_df["{self.col_name}"].isnull().any()"""
-
-class TestColumnAcceptedValues(Testing):
- def __init__(self, col_name, accepted_values, name=None, explanation="", codes="", sample_df=None):
- super().__init__(name, explanation, codes, sample_df)
- self.col_name = col_name
- if not isinstance(accepted_values, list):
- raise ValueError("Accepted values must be a list.")
- self.accepted_values = accepted_values
-
- if name is None:
- self.name = f"Test {col_name} domain"
-
- def generate_test_accepted_values_codes(self):
- self.codes = f"""def test(input_df):
- return input_df["{self.col_name}"].isin({self.accepted_values}).all()"""
-
-class TestColumnType(Testing):
- def __init__(self, col_name, expected_type, name=None, explanation="", codes="", sample_df=None):
- super().__init__(name, explanation, codes, sample_df)
- self.col_name = col_name
- self.expected_type = expected_type
-
- if name is None:
- self.name = f"Test {col_name} type"
-
- def generate_test_column_type_codes(self):
- self.codes = f"""def test(input_df):
- return input_df["{self.col_name}"].dtype == "{self.expected_type}"""
-
-class TestColumnRange(Testing):
- def __init__(self, col_name, min_value, max_value, name=None, explanation="", codes="", sample_df=None):
- super().__init__(name, explanation, codes, sample_df)
- self.col_name = col_name
- self.min_value = min_value
- self.max_value = max_value
-
- if name is None:
- self.name = f"Test {col_name} range"
-
- def generate_test_range_codes(self):
- self.codes = f"""def test(input_df):
- if not pd.api.types.is_numeric_dtype(input_df["{self.col_name}"]):
- return False
- return input_df["{self.col_name}"].between({self.min_value}, {self.max_value}).all()"""
-
-
-class TestColumnRegex(Testing):
- def __init__(self, col_name, regex_pattern, name=None, explanation="", codes="", sample_df=None):
- super().__init__(name, explanation, codes, sample_df)
- self.col_name = col_name
- self.regex_pattern = regex_pattern
-
- if name is None:
- self.name = f"Test {col_name} regex pattern"
-
- def generate_test_regex_codes(self):
- self.codes = f"""def test(input_df):
- return input_df["{self.col_name}"].str.match("{self.regex_pattern}").all()"""
-
-
-
-class TransformationStep:
-
- def __init__(self, name = None, explanation="", codes="", sample_df=None):
- if name is None:
- self.name = "Transformation Task"
- else:
- self.name = name
-
- self.codes = codes
- self.explanation = explanation
-
- if sample_df is not None:
- try:
- self.sample_df = sample_df.copy()
- except:
- self.sample_df = sample_df
-
- self.reason = ""
-
-
- def verify_infput(self, df):
- if not isinstance(df, pd.DataFrame):
- raise ValueError(f"Input is not a pandas dataframe. It is {type(df)}.")
-
- def verify_output(self, df):
- if not isinstance(df, pd.DataFrame):
- raise ValueError(f"Output is not a pandas dataframe. It is {type(df)}.")
-
-
- def generate_codes(self, explanation=None):
- if explanation is None:
- explanation = self.explanation
-
- template = f"""Transformation task: Given input df, write python codes that transform and ouput df.
-===
-Input Df:
-{self.sample_df.to_csv()}
-===
-Transformation Requirement
-{explanation}
-
-Do the following:
-1. First reason about how to transform
-2. Then fill in the python function, with detailed comments.
-DONT change the function name, first line and the return clause.
-
-{{
- "reason": "To transform, we need to ...",
- "codes": "def transform(input_df):\\n output_df = input_df.copy()\\n ...\\n return output_df"
-}}
-"""
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
- self.reason = json_code["reason"]
- self.codes = json_code["codes"]
-
-
- def postprocessing(self, df):
- return df
-
- def run_codes(self, dfs, codes = None):
-
- if codes is None:
- codes = self.codes
-
- if isinstance(dfs, list):
- df = dfs[0]
- else:
- df = dfs
-
- input_df = df.copy()
-
-
- try:
- if 'transform' in globals():
- del globals()['transform']
- exec(codes, globals())
- temp_target_df = transform(input_df)
- temp_target_df = self.postprocessing(temp_target_df)
- self.verify_output(temp_target_df)
- return temp_target_df
- except Exception:
- detailed_error_info = get_detailed_error_info()
- return detailed_error_info
-
-
-
-
- def edit_widget(self, callbackfunc=None):
-
- print("\033[91mRemember to save after editing!\033[0m")
-
- explanation_label = widgets.Label('Task:')
- explanation_text = widgets.Textarea(
- value = self.explanation if self.explanation != "" else "Transform ... For example ...",
- layout=widgets.Layout(width='95%', height='200px')
- )
- submit_button = widgets.Button(description="Submit Task")
- submit_spinner = widgets.HTML()
-
- def on_submit_clicked(b):
- print("Generating codes...")
-
- submit_spinner.value = spinner_value
- self.generate_codes(explanation = explanation_text.value)
-
- codes_text.value = self.codes
- reason_label.value = self.reason
- submit_spinner.value = ""
- print("Done")
-
- submit_button.on_click(on_submit_clicked)
- submit_box = widgets.HBox([submit_button, submit_spinner])
-
- codes_label = widgets.Label('Codes:')
- codes_text = widgets.Textarea(
- value=self.codes,
- layout=widgets.Layout(width='95%', height='200px')
- )
- run_button = widgets.Button(description="Run Codes")
- run_spinner = widgets.HTML()
-
- reason_label = widgets.Label(layout=Layout(width='100%', overflow='auto', white_space='pre-wrap'))
- output_label = widgets.Label()
-
- def on_run_clicked(b):
- print("Running codes...")
- run_spinner.value = spinner_value
- output = self.run_codes(dfs = self.sample_df, codes = codes_text.value)
- if isinstance(output, str):
- error_label.value = "" + output.replace("\n", " ") + ""
- output_df_widget.value = ""
- else:
- error_label.value = ""
- output_df_widget.value = output.to_html(border=0)
- run_spinner.value = ""
- print("Done")
-
-
- run_button.on_click(on_run_clicked)
- run_box = widgets.HBox([run_button, run_spinner])
-
-
- panel_layout = Layout(width='400px')
-
- left_panel = widgets.VBox([explanation_label, explanation_text, submit_box], layout=panel_layout)
- right_panel = widgets.VBox([codes_label, codes_text, run_box, output_label], layout=panel_layout)
- display(widgets.HBox([left_panel, right_panel]))
- display(reason_label)
-
- input_df_label = widgets.Label('Input Table:')
- display(input_df_label)
- display(HTML(self.sample_df.to_html(border=0)))
-
- output_df_label = widgets.Label('Output Table:')
- output_df_widget = widgets.HTML(
- value='',
- placeholder='Output DataFrame will be shown here',
- description='',
- )
-
- error_label = widgets.HTML(layout=Layout(overflow='auto'))
-
- save_button = widgets.Button(description="Save the Step")
-
- def on_save_clicked(b):
- print("Saving ...")
- self.codes = codes_text.value
- self.explanation = explanation_text.value
- self.reason = reason_label.value
- callbackfunc(self)
- print("Done")
-
- save_button.on_click(on_save_clicked)
-
- display(widgets.VBox([
- output_df_label,
- output_df_widget,
- error_label]))
-
- display(save_button)
-
- if self.codes != "":
- on_run_clicked(run_button)
-
- def get_sample_output(self):
- result = self.run_codes(self.sample_df)
- if isinstance(result, str):
- raise ValueError(f"The codes for step {self.name} is not correct: {result} \n Please edit the codes.")
- elif isinstance(result, pd.DataFrame):
- return result
- else:
- raise ValueError("Output is neither a string nor a pandas dataframe.")
-
- def display(self):
- print(f"{BOLD}{self.name}{END}: {self.explanation}")
- display(HTML(f""))
- print(f"{BOLD}Codes{END}:")
- print(highlight(self.codes, PythonLexer(), Terminal256Formatter()))
-
- if hasattr(self, 'sample_df'):
- display(HTML(f"Example Input: {self.sample_df.to_html()}"))
-
- if not hasattr(self, 'output_sample_df'):
- self.output_sample_df = self.run_codes(self.sample_df)
-
- display(HTML(f"Example Output: {self.output_sample_df.to_html()}"))
-
- def __repr__(self):
- self.display()
- return ""
-
- def rename_based_on_explanation(self):
- title, message = give_title(self.explanation)
- self.name = title
-
- def get_codes(self, target_id=0, source_ids=[]):
- source_ids_str = ", ".join([f"df_{source_id}" for source_id in source_ids])
- return self.codes + "\n\n" + f"df_{target_id} = transform({source_ids_str})\n\n"
-
-
-
-class GeoAggregationStep(TransformationStep):
-
- def __init__(self, shape_data, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.shape_data = shape_data
- self.agg_cols = [shape_data.x_att, shape_data.y_att]
-
- self.explanation = f"""Aggregate {self.agg_cols} by ..."""
-
- if 'name' not in kwargs:
- self.name = f"Aggregation"
-
- def verify_input(self, df):
- if not set(self.agg_cols).issubset(df.columns):
- raise ValueError(f"Columns {self.agg_cols} are not in the input dataframe.")
-
- def postprocessing(self, df):
- if not isinstance(df, pd.DataFrame):
- raise ValueError("Output is not a pandas dataframe.")
-
- if not set(self.agg_cols).issubset(df.columns):
- raise ValueError(f"Columns {self.agg_cols} are not in the output dataframe.")
-
- if df.index.name is not None:
- df = df.reset_index()
- return df
-
- def generate_codes(self, explanation=None):
- if explanation is None:
- explanation = self.explanation
-
- template = f"""Transformation task: Given input df, write python codes that aggregate and ouput df.
-===
-Input Df:
-{self.sample_df.df[:2].to_csv()}
-===
-Transformation Requirement:
-Aggregate {self.agg_cols}.
-{explanation}
-
-Do the following:
-1. First reason about how to transform. The final output df should group by only attributes: {self.agg_cols}
-2. Then fill in the python function, with detailed comments.
-DONT change the function name, first line and the return clause.
-
-{{
- "reason": "To transform, we need to ...",
- "codes": "def transform(input_df):\\n output_df = input_df.copy()\\n ...\\n return output_df"
-}}
-"""
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
- self.reason = json_code["reason"]
- self.codes = json_code["codes"]
-
-
- def run_codes(self, dfs, codes = None):
-
- if codes is None:
- codes = self.codes
-
- if isinstance(dfs, list):
- shape_data = dfs[0]
- else:
- shape_data = dfs
-
- data = shape_data.get_data()
- meta = shape_data.meta.copy()
- meta['aggregated'] = True
-
-
- try:
- if 'transform' in globals():
- del globals()['transform']
- exec(codes, globals())
- temp_output_df = transform(data)
- temp_output_df = self.postprocessing(temp_output_df)
-
- return shape_manufacturing(shape_data = temp_output_df, x_att=self.agg_cols[0], y_att=self.agg_cols[1], meta=meta)
-
- except Exception:
- detailed_error_info = get_detailed_error_info()
- return detailed_error_info
-
-
-class MultiDFTransformationStep(TransformationStep):
- def __init__(self, name=None, explanation="", codes="", sample_dfs=None):
- super().__init__(name, explanation, codes)
- if sample_dfs is not None:
- self.sample_dfs = [df.copy() for df in sample_dfs]
- else:
- self.sample_dfs = []
-
- def run_codes(self, dfs):
- input_dfs = [df.copy() for df in dfs]
- try:
- if 'transform' in globals():
- del globals()['transform']
- exec(self.codes, globals())
- temp_target_df = transform(*input_dfs)
- return temp_target_df
- except Exception:
- detailed_error_info = get_detailed_error_info()
- return detailed_error_info
-
-
-
- def display(self):
- display(HTML(f"{self.name}: {self.explanation} Codes:"))
- print(highlight(self.codes, PythonLexer(), Terminal256Formatter()))
-
- if hasattr(self, 'sample_dfs'):
- for i, df in enumerate(self.sample_dfs):
- display(HTML(f"Example Input {i+1}: {df.to_html()}"))
-
- if not hasattr(self, 'output_sample_dfs'):
- self.output_sample_dfs = self.run_codes(self.sample_dfs)
-
- for i, df in enumerate(self.output_sample_dfs):
- display(HTML(f"Example Output {i+1}: {df.to_html()}"))
-
- def generate_codes(self):
- raise NotImplementedError("generate_codes() is not implemented for MultiDFTransformationStep.")
-
-
-class ConcatenateHorizontalStep(MultiDFTransformationStep):
- def __init__(self, name=None, explanation="", codes="", sample_dfs=None):
- super().__init__(name, explanation, codes, sample_dfs)
- self.generate_concatenate_horizontal_codes()
-
- self.explanation = """Concatenate all dataframes horizontally"""
-
- if name is None:
- self.name = "Concatenate Horizontal"
-
- def generate_concatenate_horizontal_codes(self):
- self.codes = """# Concatenate all dataframes horizontally
-def transform(*dfs):
- return pd.concat(dfs, axis=1)"""
-
-
-
-class SourceStep(TransformationStep):
- def __init__(self, doc_df, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.explanation = ""
- self.doc_df = doc_df
- self.name = doc_df.table_name
-
- def generate_codes(self):
- pass
-
- def run_codes(self, dfs):
- return self.doc_df.original_df
-
- def display(self):
- display(HTML(self.doc_df.original_df[:5].to_html()))
-
- def edit_widget(self, callbackfunc=None):
- self.display()
- print("\033[91mEdit Source File is under development!\033[0m")
-
- return_button = widgets.Button(description="Return")
-
- def on_return_clicked(b):
- clear_output(wait=True)
- callbackfunc(self)
-
- return_button.on_click(on_return_clicked)
-
- display(return_button)
-
- def get_codes(self, target_id=0, source_ids=[]):
- if self.doc_df.file_path is None:
- return f"df_{str(target_id)} = pd.read_csv(ADD_YOUR_SOURCE_FILE_PATH_HERE) \n\n"
- else:
- sep = self.doc_df.sep
- encoding = self.doc_df.encoding
- return f"df_{str(target_id)} = pd.read_csv('{self.doc_df.file_path}', sep='{sep}', encoding='{encoding}') \n\n"
-
-class RemoveMissingValueStep(TransformationStep):
-
- def __init__(self, col, reason, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.col = col
- self.explanation = f"""Remove the rows with missing values in column {col}"""
- self.generate_remove_missing_value_codes()
- self.generate_missing_value_samples()
-
- if 'name' not in kwargs:
- self.name = f"Remove NULL for {col}"
-
-
- def generate_remove_missing_value_codes(self):
- self.codes = f"""# Remove the rows with missing values in column {self.col}
-def transform(df):
- output_df = df.copy()
- output_df = output_df.dropna(subset=["{self.col}"])
- return output_df"""
-
-
- def generate_missing_value_samples(self):
- sample_df1 = self.sample_df[self.sample_df[self.col].notnull()][:2]
- sample_df2 = self.sample_df[self.sample_df[self.col].isnull()][:2]
- sample_df = pd.concat([sample_df1, sample_df2])
- self.sample_df = sample_df
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-class GeoShapeCustomTransformStep(TransformationStep):
-
- def generate_codes(self, explanation=None):
-
- data_type = self.sample_df.get_type()
-
- if data_type == "raster":
- raise NotImplementedError("generate_codes() is not implemented for Raster.")
- elif data_type == "df":
- raise NotImplementedError("generate_codes() is not implemented for df.")
-
- if explanation is None:
- explanation = self.explanation
-
- template = f"""Transformation task: Given input {data_type}, write python codes that transform and ouput {data_type}.
-===
-Input {data_type}:
-{self.sample_df.get_summary()}
-===
-Transformation Requirement
-{explanation}
-
-Do the following:
-1. First reason about how to transform
-2. Then fill in the python function, with detailed comments.
-DONT change the function name and the return clause.
-
-{{
- "reason": "To transform, we need to ...",
- "codes": "def transform(input_{data_type}):\\n ...\\n return output_{data_type}"
-}}
-"""
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- write_log(template)
- write_log("-----------------------------------")
- write_log(response['choices'][0]['message']['content'])
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
- self.reason = json_code["reason"]
- self.codes = json_code["codes"]
-
- def run_codes(self, dfs, codes = None):
-
- if codes is None:
- codes = self.codes
-
- if isinstance(dfs, list):
- shape_data = dfs[0]
- else:
- shape_data = dfs
-
- data = shape_data.get_data()
- meta = shape_data.meta
-
-
- try:
- if 'transform' in globals():
- del globals()['transform']
- exec(codes, globals())
- temp_output_df = transform(data)
- return shape_manufacturing(shape_data = temp_output_df, meta = meta)
- except Exception:
- detailed_error_info = get_detailed_error_info()
- return detailed_error_info
-
-class GeoShapeStep(TransformationStep):
- def generate_codes(self):
- self.codes = f"""# Reproject to {self.target_crs}
-NOT IMPLEMENTED YET!"""
-
- def verify_input(self, shape_data):
- if not isinstance(shape_data, ShapeData):
- raise TypeError("shape_data must be of type ShapeData.")
-
- def verify_output(self, shape_data):
- if not isinstance(shape_data, ShapeData):
- raise TypeError("shape_data must be of type ShapeData.")
-
- def edit_widget(self, callbackfunc=None):
- print("🚧 under development")
- self.display()
-
- return_button = widgets.Button(description="Return")
-
- def on_return_clicked(b):
- clear_output(wait=True)
- callbackfunc(self)
-
- return_button.on_click(on_return_clicked)
-
- display(return_button)
-
- def get_sample_output(self):
- result = self.run_codes(self.sample_df)
-
- def display(self):
- if hasattr(self, 'sample_df'):
- display(HTML(f"Example Input:"))
- self.sample_df.__repr__()
-
- if not hasattr(self, 'output_sample_df'):
- self.output_sample_df = self.run_codes(self.sample_df)
-
- display(HTML(f"Example Output:"))
- self.output_sample_df.__repr__()
-
-
-class ReprojectStep(GeoShapeStep):
-
- def __init__(self, target_crs=None, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.target_crs = target_crs
-
- if 'name' not in kwargs:
- self.name = f"Reproject to {target_crs}"
-
- def run_codes(self, dfs):
- if isinstance(dfs, list):
- shape_data = dfs[0]
- else:
- shape_data = dfs
- return shape_data.project_to_target_crs(target_crs=self.target_crs)
-
-class ResampleStep(GeoShapeStep):
- def __init__(self, geo_transform=None, resolution=None, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.geo_transform = geo_transform
- self.resolution = resolution
-
- if 'name' not in kwargs:
- self.name = f"Rasample"
-
- if geo_transform is not None:
- self.name += f" to {geo_transform}"
-
- if resolution is not None:
- self.name += f" to resolution {resolution}"
-
- def run_codes(self, dfs):
- if isinstance(dfs, list):
- shape_data = dfs[0]
- else:
- shape_data = dfs
- return shape_data.resample_to_target_transform(geo_transform=self.geo_transform, resolution=self.resolution)
-
-class NumpyToDfStep(GeoShapeStep):
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- if 'name' not in kwargs:
- self.name = f"NumPy to Pandas DataFrame"
-
- def run_codes(self, dfs):
- if isinstance(dfs, list):
- shape_data = dfs[0]
- else:
- shape_data = dfs
- return shape_data.to_df()
-
-
-class SourceShapeStep(GeoShapeStep):
- def __init__(self, shape_data, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.explanation = ""
- self.shape_data = shape_data
-
- if "table_name" in shape_data.meta:
- self.name = shape_data.meta["table_name"]
- else:
- self.name = "Source Geographical Shape"
-
- def generate_codes(self):
- pass
-
- def run_codes(self, dfs):
- return self.shape_data
-
- def display(self):
- self.shape_data.display()
-
- def edit_widget(self, callbackfunc=None):
- self.display()
- print("\033[91mEdit Source File is under development!\033[0m")
-
- return_button = widgets.Button(description="Return")
-
- def on_return_clicked(b):
- clear_output(wait=True)
- callbackfunc(self)
-
- return_button.on_click(on_return_clicked)
-
- display(return_button)
-
- def get_codes(self, target_id=0, source_ids=[]):
- if hasattr(self, 'gdf'):
- return f"df_{str(target_id)} = gpd.read_file(ADD_YOUR_SOURCE_FILE_PATH_HERE) \n\n"
- elif hasattr(self, 'raster_path'):
- return f"df_{str(target_id)} = rasterio.open('{self.raster_path}').read() \n\n"
- elif hasattr(self, 'np_array'):
- return f"df_{str(target_id)} = np.load(ADD_YOUR_SOURCE_FILE_PATH_HERE) \n\n"
- elif hasattr(self, 'df'):
- return f"df_{str(target_id)} = pd.read_csv(ADD_YOUR_SOURCE_FILE_PATH_HERE) \n\n"
-
-
-
-class RemoveColumnsStep(TransformationStep):
-
- def __init__(self, col_indices, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.col_indices = col_indices
-
- if self.explanation == "":
- self.explanation = f"""Remove columns"""
-
- self.generate_remove_columns_codes()
- self.generate_remove_columns_samples()
-
- if 'name' not in kwargs:
- self.name = f"Remove Columns"
-
-
- def generate_remove_columns_codes(self):
- if isinstance(self.col_indices, list):
- cols_indices_str = str(self.col_indices)
- else:
- cols_indices_str = f"[{self.col_indices}]"
-
- self.codes = f"""# Remove columns by indices {cols_indices_str}
-def transform(df):
- # Create a boolean mask for all columns
- mask = np.full(df.shape[1], False)
- # Set True for columns to be dropped
- mask[{cols_indices_str}] = True
- # Drop columns based on the mask
- output_df = df.loc[:, ~mask]
- return output_df"""
-
- def generate_remove_columns_samples(self):
- sample_df = self.sample_df[:4]
- self.sample_df = sample_df
-
-
-
-
-class RemoveRowsStep(TransformationStep):
-
- def __init__(self, rows, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.rows = rows
-
- if self.explanation == "":
- self.explanation = f"""Remove rows {rows}"""
-
- self.generate_remove_rows_codes()
- self.generate_remove_rows_samples()
-
- if 'name' not in kwargs:
- self.name = f"Remove Rows"
-
- def generate_remove_rows_codes(self):
- if isinstance(self.rows, list):
- rows_str = str(self.rows)
- else:
- rows_str = f"[{self.rows}]"
-
- self.codes = f"""# Remove rows {rows_str}
-def transform(df):
- output_df = df.drop(index={rows_str})
- return output_df"""
-
- def generate_remove_rows_samples(self):
- sample_df1 = self.sample_df[~self.sample_df.index.isin(self.rows)][:2]
- sample_df2 = self.sample_df[self.sample_df.index.isin(self.rows)][:2]
- sample_df = pd.concat([sample_df1, sample_df2])
- self.sample_df = sample_df
-
-class RemoveDuplicatesStep(TransformationStep):
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- if self.explanation == "":
- self.explanation = f"""Remove duplicate rows"""
-
- self.generate_remove_duplicates_codes()
- self.generate_remove_duplicates_samples()
-
- if 'name' not in kwargs:
- self.name = f"Remove Duplicates"
-
- def generate_remove_duplicates_codes(self):
- self.codes = f"""# Remove duplicate rows
-def transform(df):
- output_df = df.drop_duplicates()
- return output_df"""
-
- def generate_remove_duplicates_samples(self):
- duplicate_mask = self.sample_df.duplicated(keep=False)
- duplicated_rows = self.sample_df[duplicate_mask]
-
- sample_duplicated_rows = duplicated_rows[:4]
- self.sample_df = sample_duplicated_rows
-
-class CleanDataType(TransformationStep):
-
- def __init__(self, column_data_type_dict , *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- self.column_data_type_dict = column_data_type_dict
-
- if self.explanation == "":
- self.explanation = f"""Remove rows that don't match the data type"""
-
- self.generate_clean_data_type_codes()
-
- if 'name' not in kwargs:
- self.name = f"Clean Data Type"
-
- def generate_clean_data_type_codes(self):
- self.codes = f"""# Clean data type
-def transform(df):
- data_type_dict = {self.column_data_type_dict}
- for column_name, data_type in data_type_dict.items():
- mask = select_invalid_data_type(df, column_name, data_type)
- df = df[~mask]
- return df"""
-
-
-class ProjectionStep(TransformationStep):
-
- def __init__(self, cols, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.cols = cols
-
- all_cols = ", ".join(self.cols)
-
- self.explanation = f"""Keep only column {all_cols}"""
- self.generate_projection_codes()
-
- if hasattr(self, 'sample_df'):
- self.generate_projection_samples()
-
- if 'name' not in kwargs:
- self.name = f"Project {all_cols}"
-
- def generate_projection_codes(self):
- all_cols = ', '.join([f'"{col}"' for col in self.cols])
-
- self.codes = f"""# Keep only specified columns
-def transform(df):
- output_df = df.copy()
- output_df = output_df[[{all_cols}]]
- return output_df"""
-
- def generate_projection_samples(self):
- sample_df = self.sample_df[:4]
- self.sample_df = sample_df
-
-
-
-class RegexTransformationStep(TransformationStep):
-
- def __init__(self, col, unique_values, reason, *args, **kwargs):
-
- super().__init__(*args, **kwargs)
-
- if 'name' not in kwargs:
- self.name = f"Clean {col}"
-
- self.explanation = f"""Remove the unusual values from column {col}: {reason}"""
-
- self.col = col
- self.unique_values = unique_values
- self.reason = reason
-
- def generate_regex_codes(self, explanation=None):
- if not hasattr(self, 'labeled_result'):
- raise ValueError("Please run the edit_widget() function first.")
-
- unusual_values = self.labeled_result["unusual_values"]
- normal_values = self.labeled_result["normal_values"]
-
-
- self.sample_df[self.col] = self.sample_df[self.col].astype(str)
-
- sample_df1 = self.sample_df[self.sample_df[self.col].isin(normal_values)][:2]
- sample_df2 = self.sample_df[self.sample_df[self.col].isin(unusual_values)][:2]
- sample_df = pd.concat([sample_df1, sample_df2])
- self.sample_df = sample_df
-
- print("Generating codes...")
-
- progress = show_progress(1)
- try:
- regex_result = find_regex_pattern(self.col, self.labeled_result)
- except Exception as e:
- write_log(e)
- regex_result = {"exists_regex": False}
-
- progress.value += 1
-
- if regex_result["exists_regex"]:
- regex_pattern = regex_result["regex"]
- self.codes = f"""# Clean column {self.col} by removing the unusual rows that don't match the regex pattern
-def transform(df):
- output_df = df.copy()
-
- # Transform the column to string type
- output_df["{self.col}"] = output_df["{self.col}"].astype(str)
-
- output_df = output_df[output_df["{self.col}"].str.match(r"{regex_pattern}")]
- return output_df"""
-
- self.explanation = f"""Remove the unusual values from column {self.col} by removing the rows that don't match the regex pattern: {regex_pattern}"""
-
- else:
- unusual_values = self.labeled_result["unusual_values"]
- self.codes = f"""# Clean column {self.col} by removing the unusual rows with unusual values
-def transform(df):
- output_df = df.copy()
-
- # Transform the column to string type
- output_df["{self.col}"] = output_df["{self.col}"].astype(str)
-
- output_df = output_df[~output_df["{self.col}"].isin({unusual_values})]
- return output_df"""
-
- self.explanation = f"""Remove the unusual values from column {self.col} by removing the rows with unusual values: {unusual_values}"""
-
- def parent_edit_widget(self, callbackfunc=None):
- super().edit_widget(callbackfunc)
-
-
- def edit_widget(self, callbackfunc=None):
-
- print("Analyzing the unusual values...")
-
- if not hasattr(self, 'labeled_result'):
- progress = show_progress(1)
- result = classify_unusual_values(self.col, self.unique_values, self.reason)
- self.labeled_result = result
- progress.value += 1
-
- def create_tabular_widgets_toggle(result):
- grid_items = []
-
- grid_items.append(widgets.Label(value=''))
- grid_items.append(widgets.Label(value='Keep (Normal)'))
- grid_items.append(widgets.Label(value='Remove (Unusual)'))
-
- checkbox_widgets = {}
-
- for email in result['normal_values'] + result['unusual_values']:
- is_normal = email in result['normal_values']
-
- label = widgets.Label(value=email)
- grid_items.append(label)
-
- keep_checkbox = widgets.Checkbox(value=is_normal)
- grid_items.append(keep_checkbox)
-
- remove_checkbox = widgets.Checkbox(value=not is_normal)
- grid_items.append(remove_checkbox)
-
- checkbox_widgets[email] = (keep_checkbox, remove_checkbox)
-
- keep_checkbox.observe(lambda change, email=email: on_checkbox_toggle(change, email, 'keep'), names='value')
- remove_checkbox.observe(lambda change, email=email: on_checkbox_toggle(change, email, 'remove'), names='value')
-
- def on_checkbox_toggle(change, email, action):
- if change['new']:
- if action == 'keep':
- checkbox_widgets[email][1].value = False
- else:
- checkbox_widgets[email][0].value = False
-
- grid_layout = widgets.Layout(grid_template_columns="repeat(3, 33%)",
- align_items='center')
- grid_box = widgets.GridBox(children=grid_items, layout=grid_layout)
-
- submit_button = widgets.Button(description="Submit")
-
- def on_submit_button_clicked(b):
- clear_output(wait=True)
-
- updated_result = {'unusual_values': [], 'normal_values': []}
- for email, (keep_checkbox, remove_checkbox) in checkbox_widgets.items():
- if keep_checkbox.value:
- updated_result['normal_values'].append(email)
- else:
- updated_result['unusual_values'].append(email)
-
- self.labeled_result = updated_result
-
- self.generate_regex_codes()
-
- print("Done")
-
- self.parent_edit_widget(callbackfunc)
-
-
- submit_button.on_click(on_submit_button_clicked)
-
- display(grid_box, submit_button)
-
- print("Please verify the unusual values...")
-
- create_tabular_widgets_toggle(self.labeled_result)
-
-
-
-class ColumnRename(TransformationStep):
-
- def __init__(self, rename_map, *args, **kwargs):
-
- super().__init__(*args, **kwargs)
-
- if 'name' not in kwargs:
- self.name = "Rename Columns"
-
- self.rename_map = rename_map
- self.explanation = f"""Rename columns in the DataFrame based on: \n"""
-
- valid_renaming = {}
-
- for old_name, new_name in self.rename_map.items():
- if old_name != new_name:
- valid_renaming[old_name] = new_name
-
- rename_map_str = "{\n"
- rename_map_str += "\n".join([f" '{old}': '{new}'," for old, new in valid_renaming.items()])
- rename_map_str += "\n }"
-
- self.explanation += rename_map_str
-
- self.generate_rename_codes()
-
- def verify_input(self, df):
- super().verify_input(df)
-
- if len(self.rename_map) != len(set(self.rename_map.keys())):
- raise ValueError("Some old column names are duplicated.")
-
- if len(self.rename_map) != len(set(self.rename_map.values())):
- raise ValueError("Some new column names are duplicated.")
-
- for old_name in self.rename_map.keys():
- if old_name not in df.columns:
- raise ValueError(f"Column {old_name} does not exist in the DataFrame.")
-
- def generate_rename_codes(self):
-
- valid_renaming = {}
-
- for old_name, new_name in self.rename_map.items():
- if old_name != new_name:
- valid_renaming[old_name] = new_name
-
- rename_map_str = "{\n"
- rename_map_str += "\n".join([f" '{old}': '{new}'," for old, new in valid_renaming.items()])
- rename_map_str += "\n }"
-
- self.codes = f"""# Rename columns in the DataFrame based on column indices to avoid circular issues
-def transform(df):
- rename_map = {rename_map_str}
- # Create a list of the current column names
- new_column_names = list(df.columns)
-
- # Find the indices of the columns to be renamed and update their names
- for old_name, new_name in rename_map.items():
- if old_name in df.columns:
- index = df.columns.get_loc(old_name)
- new_column_names[index] = new_name
- else:
- raise ValueError(f'Column {{{{old_name}}}} not found in DataFrame')
-
- # Assign the new names to the DataFrame's columns
- df.columns = new_column_names
- return df
-"""
-
-
-
-
-
-
-
-class DocumentedDatabase:
- def __init__(self, doc_dfs):
- self.doc_dfs = doc_dfs
-
- def generate_pipeline(self):
- pipelines = []
-
- for doc_df in self.doc_dfs:
- pipelines.append(doc_df.generate_pipeline())
-
- final_pipeline = combine_pipelines(pipelines)
- return final_pipeline
-
-class DataCleaning:
- def __init__(self, doc_df):
- self.doc_df = doc_df
- self.pipeline = doc_df.generate_pipeline()
-
- def __repr__(self):
- self.display()
- return ""
-
- def run_codes(self):
- return self.pipeline.run_codes()
-
- def display(self):
-
-
- self.pipeline.display(call_back=self.display)
-
- attributes = []
- issues = {}
-
- document = self.doc_df.document
-
- if "missing_value" in document:
- issues["missing_value"] = []
- for attribute in document["missing_value"]:
- item = document["missing_value"][attribute]
- if item:
- attributes.append(attribute)
- issues["missing_value"].append((attribute, item["summary"]))
-
- if "unusual" in document:
- issues["unusual"] = []
- for attribute in document["unusual"]:
- item = document["unusual"][attribute]["summary"]
- if item["Unusualness"]:
- attributes.append(attribute)
- issues["unusual"].append((attribute, item["Examples"]))
-
- self.recommends_transformation(issues)
-
- def recommends_transformation(self, issues):
-
- def on_button_clicked_missing(b):
- clear_output(wait=True)
- self.create_remove_missing_value_step(b.attribute, b.issue)
-
- boxes = []
- for attribute, issue in issues["missing_value"]:
- label = widgets.HTML(value=f"❓ {attribute} has missing values: {issue}")
-
- button = widgets.Button(description=f"Remove them")
- button.attribute = attribute
- button.issue = issue
- button.on_click(on_button_clicked_missing)
-
- box = widgets.VBox([label, button])
- boxes.append(box)
-
- def on_button_clicked(b):
- clear_output(wait=True)
- unique_values = self.doc_df.df[b.attribute].dropna().unique()[:20]
- self.create_regex_step(b.attribute, b.issue, unique_values)
-
- for attribute, issue in issues["unusual"]:
- label = widgets.HTML(value=f"🤔 {attribute} has unusual value: {issue}")
-
- button = widgets.Button(description=f"Remove them")
- button.attribute = attribute
- button.issue = issue
- button.on_click(on_button_clicked)
-
- box = widgets.VBox([label, button])
- boxes.append(box)
-
- def on_remove_last_clicked(b):
- clear_output(wait=True)
- self.add_remove_last_step()
-
- remove_last_button = widgets.Button(description="⚠️ Remove Last")
- remove_last_button.on_click(on_remove_last_clicked)
- box = widgets.VBox([remove_last_button])
- boxes.append(box)
-
-
- label = widgets.HTML(value=f"🎲 Want to perform an ad hoc transformation?")
-
- def on_ad_hoc_clicked(b):
- clear_output(wait=True)
- self.create_ad_hoc_step()
-
- adhoc_button = widgets.Button(description="Ad hoc")
- adhoc_button.on_click(on_ad_hoc_clicked)
- box = widgets.VBox([label, adhoc_button])
- boxes.append(box)
-
-
- display(widgets.VBox(boxes))
-
-
- def add_remove_last_step(self):
- self.pipeline.remove_final_node()
- self.final_shape = self.pipeline.run_codes()
- clear_output(wait=True)
- self.display()
-
- def create_remove_missing_value_step(self, col, reason):
-
-
- sample_output = self.doc_df.df
-
- remove_missing_value_step = RemoveMissingValueStep(col=col, reason=reason, sample_df=sample_output)
-
- def callback(remove_missing_value_step):
- if remove_missing_value_step.explanation != "" or remove_missing_value_step.codes != "":
- self.pipeline.add_step_to_final(remove_missing_value_step)
- clear_output(wait=True)
- self.display()
-
- callbackfunc = callback
-
- remove_missing_value_step.edit_widget(callbackfunc=callbackfunc)
-
-
- def create_regex_step(self, col, reason, unique_values):
-
-
- sample_output = self.doc_df.df
-
- regex_step = RegexTransformationStep(col=col, unique_values=unique_values, reason=reason, sample_df=sample_output)
-
- def callback(regex_step):
- if regex_step.explanation != "" or regex_step.codes != "":
-
- self.pipeline.add_step_to_final(regex_step)
-
- clear_output(wait=True)
- self.display()
-
- callbackfunc = callback
-
- regex_step.edit_widget(callbackfunc=callbackfunc)
-
- def create_ad_hoc_step(self):
-
- final_node = self.pipeline.find_final_node()
- final_step = self.pipeline.get_step(final_node)
- sample_output = final_step.get_sample_output()
-
- add_hoc_step = TransformationStep(name="Ad hoc", sample_df=sample_output)
-
- def callback(add_hoc_step):
- if add_hoc_step.explanation != "" or add_hoc_step.codes != "":
- add_hoc_step.rename_based_on_explanation()
- self.pipeline.add_step_to_final(add_hoc_step)
- clear_output(wait=True)
- self.display()
-
- callbackfunc = callback
-
- add_hoc_step.edit_widget(callbackfunc=callbackfunc)
-
- def print_codes(self):
- self.pipeline.print_codes()
-
- def generate_pipeline(self):
- return self.pipeline
-
-def replace_nan(df):
- for col in df.columns:
- df[col] = df[col].fillna("")
- return df
-
-def embed_string(string, model_name, engine=None):
- try:
- response = call_embed(string, model_name)
-
- if openai.api_type == "bedrock":
- embeddings = response
- else:
- embeddings = response['data'][0]['embedding']
- return embeddings
- except Exception as e:
- print(f"An error occurred while embedding: {e}")
- return None
-
-def initialize_output_csv(df, output_csv_address, label='label'):
- try:
- df_output = pd.read_csv(output_csv_address)
- except FileNotFoundError:
- df_output = get_unique_labels_with_ids(df, label='label')
- df_output.to_csv(output_csv_address, index=False)
- return df_output
-
-def find_first_nan_index(df, column_name):
- """
- Finds the first index in a DataFrame column that is NaN.
-
- :param df: pandas DataFrame object
- :param column_name: String name of the column to search for NaN
- :return: The index of the first NaN value, or None if no NaN found
- """
- nan_index = df[column_name].isna().idxmax()
- if pd.isna(df[column_name][nan_index]):
- return nan_index
- else:
- return None
-
-def embed_labels(df, output_csv_address, model_name, chunk_size=1000, label='label'):
- df_output = initialize_output_csv(df, output_csv_address, label='label')
-
- start_index = find_first_nan_index(df_output, 'embedding')
-
- if start_index is None:
- print("All labels already embedded.")
- return df_output
-
- pbar = tqdm(total=len(df_output), desc="Embedding Labels", unit="label")
-
- pbar.update(start_index)
-
- for chunk_start in range(start_index, len(df_output), chunk_size):
- chunk_end = min(chunk_start + chunk_size, len(df_output))
- labels_chunk = df_output[label][chunk_start:chunk_end]
-
- for i, label_value in enumerate(labels_chunk):
- if pd.isna(df_output.at[chunk_start + i, 'embedding']):
- embeddings = embed_string(label_value, model_name)
- if embeddings is not None:
- df_output.at[chunk_start + i, 'embedding'] = embeddings
- pbar.update(1)
-
- df_output.to_csv(output_csv_address, index=False)
-
- pbar.close()
-
- print("All labels embedded and CSV updated.")
-
- return df_output
-
-def get_unique_labels(df, label='label'):
- unique_labels = df[label].dropna().unique()
- return unique_labels
-
-
-def parse_json_col(df, col='embedding'):
- df[col] = df[col].apply(ast.literal_eval)
- return df
-
-
-def load_embedding(df, model_name, label_embedding='embedding', dim=1536):
- if not isinstance(df[label_embedding].iloc[0], list):
- df = parse_json_col(df, col=label_embedding)
-
- if openai.api_type == "bedrock" and model_name == "titan-v2":
- dim=1024
-
- embeddings_array = np.array(list(df[label_embedding]), dtype=np.float32)
- index = faiss.IndexFlatL2(dim)
- index.add(embeddings_array)
- return index
-
-def adhoc_search(string, index, topk=10):
- adhoc_embed = embed_string(string)
- D, I = index.search(np.array([adhoc_embed], dtype=np.float32), topk)
- return D, I
-
-def df_search(df, index, label_embedding='embedding', topk=10):
-
- if not isinstance(df[label_embedding].iloc[0], list):
- df = parse_json_col(df, col=label_embedding)
-
- embeddings = list(df[label_embedding])
- D, I = index.search(np.array(embeddings, dtype=np.float32), topk)
- return D, I
-
-def flatten_append(df, std_df, D, I, topk=10):
- df = df.loc[df.index.repeat(topk)].reset_index(drop=True)
- df['D_values'] = D.flatten()
- df['I_values'] = I.flatten()
- df['rank'] = (df.index % topk) + 1
-
- for col in std_df.columns:
- df[col] = std_df.iloc[df['I_values'].values][col].values
-
- return df
-
-def display_matches(reference_df,
- input_df,
- I,
- exclude_columns=['label', 'index_ids', 'embedding'],
- label_col = 'label'):
-
-
- current_page = 0
-
- def create_html_content(page_no):
- label = input_df[label_col][page_no]
- results_indices = I[page_no]
- results = reference_df.iloc[results_indices][label_col]
- results = [result.replace('\n', ' ') for result in results]
- results_html = "
" + "".join(f"
{result}
" for result in results) + "
"
- html_content = f"Input Label: {df_row_to_column_value(input_df, idx=page_no, exclude_columns=exclude_columns).to_html()} Top 10 Matches:{results_html}"
- return html_content
-
- def update_html_display(page_no):
- html_display.value = create_html_content(page_no)
- page_label.value = f'Page {page_no + 1} of {len(input_df)}'
-
- def on_prev_clicked(b):
- nonlocal current_page
- if current_page > 0:
- current_page -= 1
- update_html_display(current_page)
-
- def on_next_clicked(b):
- nonlocal current_page
- if current_page < len(input_df) - 1:
- current_page += 1
- update_html_display(current_page)
-
- html_display = widgets.HTML(value=create_html_content(current_page))
-
- btn_prev = widgets.Button(description='Previous Page')
- btn_next = widgets.Button(description='Next Page')
-
- btn_prev.on_click(on_prev_clicked)
- btn_next.on_click(on_next_clicked)
-
- page_label = widgets.Label(value=f'Page {current_page + 1} of {len(input_df)}')
-
- navigation_bar = widgets.HBox([btn_prev, page_label, btn_next])
-
- display(navigation_bar, html_display)
-
-
-def check_functional_dependency(df, determinant, dependent):
- groups = df.groupby(list(determinant))[dependent].nunique()
- is_functionally_dependent = (groups == 1).all()
- return is_functionally_dependent
-
-
-def get_unique_labels_with_ids(df, label = 'label'):
- df_cleaned = df.dropna(subset=['label'])
-
- grouped = df_cleaned.groupby('label')
-
- label_id_df = grouped.apply(lambda x: x.index.tolist()).reset_index(name='index_ids')
-
- dependent_columns = []
-
- for column in df_cleaned.columns:
- if column != 'label':
- if grouped[column].nunique().eq(1).all():
- dependent_columns.append(column)
-
- if dependent_columns:
- attributes_data = grouped[dependent_columns].first().reset_index()
- label_id_df = pd.merge(label_id_df, attributes_data, on='label', how='left')
-
- label_id_df['embedding'] = None
-
- return label_id_df
-
-
-def df_row_to_column_value(df, idx=0, exclude_columns=[]):
- """
- Create a DataFrame of [column, value] pairs from a specified row index,
- excluding specified columns.
-
- Parameters:
- - df: The input DataFrame.
- - idx: The index of the row to use.
- - exclude_columns: A set or list of columns to exclude.
-
- Returns:
- - A new DataFrame with two columns: 'Column' and 'Value'.
- """
-
- df_reduced = df.drop(columns=exclude_columns, errors='ignore')
-
- if idx not in df_reduced.index:
- raise IndexError(f"Index {idx} is out of bounds for the DataFrame.")
-
- row_transposed = df_reduced.iloc[idx].transpose()
-
- new_df = pd.DataFrame(row_transposed)
- new_df.reset_index(inplace=True)
- new_df.columns = ['Column', 'Value']
-
- return new_df
-
-def extract_json_code_safe(s):
- s_stripped = s.strip()
- if (s_stripped.startswith('{') and s_stripped.endswith('}')) or \
- (s_stripped.startswith('[') and s_stripped.endswith(']')):
- return s_stripped
- return extract_json_code(s_stripped)
-
-def extract_json_code(s):
- import re
- pattern = r"```json(.*?)```"
- match = re.search(pattern, s, re.DOTALL)
- return match.group(1).strip() if match else None
-
-def compute_cluster(df, match="matches"):
- clusters = {}
-
- for idx, row in df.iterrows():
- entry = row[match]
- if 'similar_to' not in entry:
- clusters[idx] = []
-
- for idx, row in df.iterrows():
- entry = row[match]
- if 'similar_to' in entry:
- similar_to_idx = entry['similar_to']
- if similar_to_idx in clusters:
- clusters[similar_to_idx].append(idx)
- elif similar_to_idx not in clusters:
- clusters[similar_to_idx] = [idx]
- return clusters
-
-def generate_report_for_cluster(df, clusters, exclude_columns=['label', 'index_ids', 'embedding','matches'], match_col='matches'):
- middle_html = ""
-
- for i in clusters:
-
- js = clusters[i]
- middle_html += f"""
-
-
-"""
- return full_html
-
-def entity_relation_match(input_df, I, refernece_df, attributes=None, label = "label", match="matches"):
-
- if match not in input_df:
- input_df[match] = None
-
- if attributes is None:
- attributes = input_df.columns.tolist()
- attributes.remove("label")
- attributes.remove("index_ids")
- attributes.remove("embedding")
-
- for idx in range(len(input_df)):
- print(f"💪 Working on the row {idx+1} ...")
-
- if input_df[match][idx] is not None:
- continue
-
- input_desc = ""
- for attribute in attributes:
- input_desc += (attribute + ": " + input_df.iloc[idx][attribute] + "\n")
-
- refernece_desc = ""
- for i, output in enumerate(refernece_df[label].iloc[I[idx]]):
- refernece_desc += (str(i+1) + ". " + output + "\n")
-
- template = f"""Your goal is to build relations between input and reference entities.
-
-The input entity has the following attributes:
-{input_desc}
-Below are reference entities:
-{refernece_desc}
-Do the following:
-1. Read input entity attributes and guess what it is about.
-
-2. Go through each output entity. Describe what it is and reason its relation.
-For instance, given the input entity "small car":
-if the same entity then EXACT_MATCH.
- E.g., "small automobile"
-else if has assumptions that is clearly wrong then CONFLICTED_ASSUMPTION
- E.g., "big car" is wrong because input entity clearly specify size as "small"
-else if additional assumptions that can't be verified then ADDITIONAL_ASSUMPTION
- E.g., "electronic car" is additional battery assumption can't be verified
-else if it is general super class then GENERAL
- E.g., "small vehicle" and "car" are general category of "small car"
-else it is irrelavent entity then NOT_RELATED
- E.g., "cloth" is a irrelavent
-
-The list is prioritized. Choose the first one that applies."""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
-
-
- reponded_content = response['choices'][0]['message']['content']
-
- messages=[
- {"role": "user",
- "content": template
- },
- {"role": "assistant",
- "content": reponded_content
- },
- {"role": "user",
-"content": """First, criticize your answer and point out the mistakes.
-For output entities classifed as EXACT_MATCH, are they making additional, potentially incorrect, or fewer assumptions? If so, correct them as ADDITIONAL_ASSUMPTION/CONFLICTED_ASSUMPTION/GENERAL.
-For output entities classifed as GENERAL, are they making additional, potentially incorrect assumptions? If so, correct them as ADDITIONAL_ASSUMPTION/CONFLICTED_ASSUMPTION.
-For output entities classifed as NOT_RELATED, are the entities similar but only part of the properties are different? If so, correct them as CONFLICTED_ASSUMPTION.
-Next, provide your corrected answer as json:
-```json
-{
- "Input Entity Guess": "...",
- "EXACT_MATCH": {
- "entity": [...],
- "reason": "The input entity is ... which matches ..."
- },
- "CONFLICTED_ASSUMPTION": {
- "entity": [...],
- "reason": "The (what specific) details are conflicted"
- },
- "ADDITIONAL_ASSUMPTION": {
- "entity": [...],
- "reason": "The (what specific) details are not mentioned"
- },
- "GENERAL": {
- "entity": [...],
- "reason": "..."
- }
- (DON'T include NOT_RELATED!)
-}
-```
-If no matched entity, return empty entity list and reason string.
-"""},]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_var = json.loads(json_code)
- input_df.at[idx, match] = json_var
-
-def entity_relation_match_cluster(input_df, I, refernece_df, attributes=None, label = "label", match="matches", verbose=False):
-
- emdeb_searcher = EmbeddingSearcher(input_df)
-
- if match not in input_df:
- input_df[match] = None
-
-
- for idx in range(len(input_df)):
- if input_df[match][idx] is not None:
- emdeb_searcher.remove_rows(idx)
-
- if attributes is None:
- attributes = input_df.columns.tolist()
- attributes.remove("label")
- attributes.remove("index_ids")
- attributes.remove("embedding")
-
- while not emdeb_searcher.is_index_empty():
- print(f"👉 {emdeb_searcher.get_size()} rows remain...")
-
-
- idx = emdeb_searcher.get_valid_id()
-
- emdeb_searcher.remove_rows(idx)
-
- input_desc = ""
- for attribute in attributes:
- input_desc += (attribute + ": " + input_df.iloc[idx][attribute] + "\n")
-
- reference_entities = list(refernece_df[label].iloc[I[idx]])
-
- refernece_desc = ""
- for i, output in enumerate(refernece_df[label].iloc[I[idx]]):
- refernece_desc += (str(i+1) + ". " + output + "\n")
- if verbose:
- print(f"👉 Input: {input_desc}")
- print(f"👉 Reference: {refernece_desc}")
-
- json_var = entity_relation_match_one(input_desc, refernece_desc)
-
- def replace_indices_with_entities(json_var, reference_entities):
- for category in json_var:
- if isinstance(json_var[category], dict) and "entity" in json_var[category]:
- json_var[category]["entity"] = [reference_entities[int(idx) - 1] for idx in json_var[category]["entity"]]
- return json_var
-
- json_var = replace_indices_with_entities(json_var, reference_entities)
-
- if verbose:
- print(f"👉 Match: {json.dumps(json_var, indent=4)}")
-
- all_indicies = []
-
-
-
- related_rows = emdeb_searcher.search_by_row_index(idx, k=30)
-
- if len(related_rows) > 0:
-
-
-
-
-
- entity_desc = json_var["Summary of Relations"]
- refernece_desc = related_rows[attributes].reset_index(drop=True).to_csv(quoting=2)
-
- json_var2 = find_relation_satisfy_description(entity_desc=entity_desc,
- related_rows_desc_str=refernece_desc)
-
- indicies = json_var2["indices"]
-
- ids = []
- for index in indicies:
- ids.append(list(related_rows.index)[index])
-
- all_indicies += ids
- emdeb_searcher.remove_rows(ids)
-
-
- json_var3 = {"similar_to": idx}
- for i, index in enumerate(all_indicies):
- input_df.at[index, match] = json_var3
-
- input_df.at[idx, match] = json_var
-
-def generate_html_from_json_entity(json_var):
- html_output = f"
🤓 This input entity is about:{json_var['Input Entity Guess']}
"
- if json_var['EXACT_MATCH']['entity']:
- html_output += "
😀 We find exactly matched entities:
"
- html_output += "
"
- for entity in json_var['EXACT_MATCH']['entity']:
- html_output += f"
{entity}
"
- html_output += "
"
- if json_var['EXACT_MATCH']['reason']:
- html_output += f"
-
'
-
- return html_str
-
-
-def get_tree_javascript(data):
- def process_data(data):
- keys = list(data.keys())
- if len(keys) == 0:
- return {}
-
- key = keys[0]
- main_data = data[key]
-
- children = []
- for sub_key, sub_value in main_data.items():
- if isinstance(sub_value, list):
- items_str = ','.join(map(str, sub_value))
-
- final_str = '[' + items_str + ']'
-
- leaf_node = [{"name": final_str}]
- children.append({"name": sub_key, "children": leaf_node})
- else:
- children.append(process_data({sub_key: sub_value}))
-
- return {"name": key, "children": children}
-
-
-
- def count_leaf_nodes(node):
- """
- Count the number of leaf nodes in a tree represented by nested dictionaries.
-
- Args:
- node (dict): The tree node, which could be the root of the tree or any sub-node within the tree.
-
- Returns:
- int: The total number of leaf nodes in the tree from the given node down.
- """
- if not isinstance(node, dict) or 'children' not in node or not node['children']:
- return 1
-
- leaf_count = 0
- for child in node['children']:
- leaf_count += count_leaf_nodes(child)
-
- return leaf_count
-
- def max_path_length(data, ch_size, gap_size):
- def calculate_length(path):
- return sum(len(step) for step in path) * ch_size + (len(path) - 1) * gap_size
-
- def traverse(node, path=[]):
- if isinstance(node, dict):
- for key, value in node.items():
- yield from traverse(value, path + [key])
- elif isinstance(node, list):
- for item in node:
- yield calculate_length(path + [item])
-
- return max(traverse(data))
-
- processed_data = process_data(data)
-
- number_of_leaves = count_leaf_nodes(processed_data)
-
- javascript = """
- const structuredData = {{data}};
-
- const treeLayout = d3.tree().size([{{height1}}, {{width1}}]);
-
- const root = d3.hierarchy(structuredData);
- root.x0 = 0;
- root.y0 = 0;
-
- const svg = d3.select('#tree').append('svg')
- .attr('width', {{width2}})
- .attr('height', {{height2}});
-
- const g = svg.append('g')
- .attr('transform', 'translate(50,50)');
-
- update(root);
-
- function update(source) {
- const duration = 750;
-
- // Re-compute the tree layout
- treeLayout(root);
-
- // Determine the leftmost node's y-position to adjust the tree positioning
- let leftmostNode = root;
- root.each(d => {
- if (d.y < leftmostNode.y) {
- leftmostNode = d;
- }
- });
- const shiftX = -leftmostNode.y + 50;
-
- const nodes = root.descendants();
- const links = root.links();
-
- const node = g.selectAll('.node')
- .data(nodes, d => d.id || (d.id = Math.random()));
-
- const nodeEnter = node.enter().append('g')
- .attr('class', 'node')
- .attr('transform', d => `translate(${source.y0 + shiftX},${source.x0})`)
- .on('click', click);
-
- nodeEnter.append('circle')
- .attr('r', 1e-6)
- .style('fill', d => d._children ? 'lightgreen' : '#fff');
-
- nodeEnter.append('text')
- .attr('dy', '0.35em')
- .attr('x', d => d.children || d._children ? -13 : 13)
- .style('text-anchor', d => d.children || d._children ? 'end' : 'start')
- .text(d => d.data.name);
-
- const nodeUpdate = nodeEnter.merge(node);
-
- nodeUpdate.transition()
- .duration(duration)
- .attr('transform', d => `translate(${d.y + shiftX},${d.x})`);
-
- nodeUpdate.select('circle')
- .attr('r', 5)
- .style('fill', d => d._children ? 'lightgreen' : '#fff');
-
- const nodeExit = node.exit().transition()
- .duration(duration)
- .attr('transform', d => `translate(${source.y + shiftX},${source.x})`)
- .remove();
-
- nodeExit.select('circle')
- .attr('r', 1e-6);
-
- nodeExit.select('text')
- .style('fill-opacity', 1e-6);
-
- const link = g.selectAll('.link')
- .data(links, d => d.target.id);
-
- const linkEnter = link.enter().insert('path', 'g')
- .attr('class', 'link')
- .attr('d', d => {
- const o = { x: source.x0, y: source.y0 };
- return diagonal(o, o);
- });
-
- const linkUpdate = linkEnter.merge(link);
-
- linkUpdate.transition()
- .duration(duration)
- .attr('d', d => diagonal(d.source, d.target));
-
- link.exit().transition()
- .duration(duration)
- .attr('d', d => {
- const o = { x: source.x, y: source.y };
- return diagonal(o, o);
- })
- .remove();
-
-
- // Update the positional attributes of each node for the next transition.
- nodes.forEach(d => {
- d.x0 = d.x;
- d.y0 = d.y;
- });
-
- function diagonal(s, d) {
- const path = `M ${s.y + shiftX} ${s.x}
- C ${(s.y + d.y) / 2 + shiftX} ${s.x},
- ${(s.y + d.y) / 2 + shiftX} ${d.x},
- ${d.y + shiftX} ${d.x}`;
-
- return path;
- }
-
- }
-
- function click(event, d) {
- if (d.children) {
- d._children = d.children;
- d.children = null;
- } else {
- d.children = d._children;
- d._children = null;
- }
- update(d);
- }"""
-
- processed_data_str = json.dumps(processed_data, indent=4)
-
- height1 = number_of_leaves * 30
- height2 = number_of_leaves * 35 + 50
-
- width1 = max_path_length(data, 5, 20)
- width2 = max_path_length(data, 10, 50) + 150
-
- html_content_updated = javascript.replace('{{data}}', processed_data_str)
- html_content_updated = html_content_updated.replace('{{height1}}', str(height1))
- html_content_updated = html_content_updated.replace('{{height2}}', str(height2))
- html_content_updated = html_content_updated.replace('{{width1}}', str(width1))
- html_content_updated = html_content_updated.replace('{{width2}}', str(width2))
-
- return html_content_updated, width2, height2
-
-
-def get_tree_html(data):
- javascript_content, width2, height2 = get_tree_javascript(data)
-
- html_content = f"""
-
-
-"""
-
- return html_content, width2, height2
-
-
-
-
-def get_rename(meanings):
-
- column_meanings = '\n'.join([f" \"{m['column']}\": {m['meaning']}" for m in meanings])
-
- template = f"""{column_meanings}
-
- Analyze the column names. The final result as a json file:
-
- ```json
- [{{
- "column": "column_name" (case sensitive),
- "ambiguous": Given the meaning, is the column name ambiguous? Note that it's fine for the name to be domain specific.
- "rename": "" (empty string if not ambiguous)
- }},...]```"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- assistant_message = response['choices'][0]['message']
- messages.append(assistant_message)
-
- for message in messages:
- write_log(message['content'])
- write_log("---------------------")
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
- def verify_rename(meanings, json_code):
-
- meanings_columns = set([m['column'] for m in meanings])
- json_code_columns = set([j['column'] for j in json_code])
- if meanings_columns != json_code_columns:
- raise ValueError(f"The set of column names in the meanings and json_code are different: {meanings_columns} vs {json_code_columns}")
-
- verify_rename(meanings, json_code)
-
- return json_code, messages
-
-
-def topological_sort(nodes, edges):
- in_degree = {i: 0 for i in range(len(nodes))}
- for targets in edges.values():
- for target in targets:
- in_degree[target] += 1
-
- queue = [i for i in range(len(nodes)) if in_degree[i] == 0]
-
- order = []
-
- while queue:
- node_idx = queue.pop(0)
- order.append(node_idx)
-
- for target_idx in edges.get(node_idx, []):
- in_degree[target_idx] -= 1
- if in_degree[target_idx] == 0:
- queue.append(target_idx)
-
- if len(order) != len(nodes):
- return "Cycle detected in the graph, topological sort not possible."
-
- return order
-
-
-
-def show_progress(max_value=1, value=1):
- progress = widgets.IntProgress(
- value=value,
- min=0,
- max=max_value+1,
- step=1,
- description='',
- bar_style='',
- orientation='horizontal'
- )
-
- display(progress)
- return progress
-
-
-def give_title(task):
- template = f"""Rename the below task into a title that's as short as possible:
-{task}"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- title = response['choices'][0]['message']['content']
- assistant_message = response['choices'][0]['message']
- messages.append(assistant_message)
-
- for message in messages:
- write_log(message['content'])
- write_log("---------------------")
-
- return title, messages
-
-
-def classify_unusual_values(col, unique_values, reason):
-
- template = f"""Issue: In the '{col}' column, some values are unusual.
-Data Values: {unique_values}
-Reason: {reason}
-Goal: Classify unusual and normal values.
-
-Return the results in json format:
-```json
-{{
- "reasoning": "The unusual values are ...",
- "unusual_values": ["unusual value",...], (could be empty)
- "normal_values": ["normal value",...] (could be empty)
-}}```"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
- assistant_message = response['choices'][0]['message']
- messages.append(assistant_message)
-
- for message in messages:
- write_log(message['content'])
- write_log("---------------------")
-
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
-
- def verify_result(json_code, unique_values):
-
- unique_values = [str(v) for v in unique_values]
-
- assert json_code['reasoning'] is not None
- assert json_code['unusual_values'] is not None
- assert json_code['normal_values'] is not None
-
- assert isinstance(json_code['reasoning'], str)
- assert isinstance(json_code['unusual_values'], list)
- assert isinstance(json_code['normal_values'], list)
-
- assert len(set(json_code['unusual_values']).intersection(set(json_code['normal_values']))) == 0
-
- json_codes_all = json_code['unusual_values'] + json_code['normal_values']
- json_codes_all = [str(v) for v in json_codes_all]
-
- assert len(set(unique_values).difference(set(json_codes_all))) == 0
-
- verify_result(json_code, unique_values)
-
- return json_code
-
-def find_regex_pattern(col, result):
-
- template = f"""Issue: In the '{col}' column, some values are unusual.
-Unusual Values: {result['unusual_values']}
-Normal Values: {result['normal_values']}
-
-Goal: Identify if there is regular expression that
-(1) matches all the normal values, and
-(2) does not match any of the unusual values.
-
-There could be multiple regex patterns that satisfy the goal.
-Please understand the meaning of the column, and find one general regex pattern.
-
-Return the results in json format:
-```json
-{{
-"reasoning": "The patterns in the normal values are ...",
-"exists_regex": true/false,
-"regex": "..."
-}}```"""
-
- messages = [{"role": "user", "content": template}]
-
- response = call_llm_chat(messages, temperature=0.1, top_p=0.1)
-
-
- assistant_message = response['choices'][0]['message']
- messages.append(assistant_message)
-
- for message in messages:
- write_log(message['content'])
- write_log("---------------------")
-
- json_code = extract_json_code_safe(response['choices'][0]['message']['content'])
- json_code = replace_newline(json_code)
- json_code = json.loads(json_code)
-
-
-
- def verify_json(json_code):
- fields = ['reasoning', 'exists_regex', 'regex']
- for field in fields:
- if field not in json_code:
- raise ValueError(f"Missing field '{field}' in json code.")
-
- if not isinstance(json_code['reasoning'], str):
- raise ValueError(f"Field 'reasoning' should be string.")
- if not isinstance(json_code['exists_regex'], bool):
- raise ValueError(f"Field 'exists_regex' should be boolean.")
- if not isinstance(json_code['regex'], str):
- raise ValueError(f"Field 'regex' should be string.")
-
- verify_json(json_code)
-
- if json_code['exists_regex']:
-
- def verify_regex(regex, normal_values, unusual_values):
- regex = re.compile(regex)
- for value in normal_values:
- if not regex.match(value):
- raise ValueError(f"Regex does not match normal value: {value}")
- for value in unusual_values:
- if regex.match(value):
- raise ValueError(f"Regex matches unusual value: {value}")
-
- verify_regex(json_code["regex"], result['normal_values'], result['unusual_values'])
-
- return json_code
-
-
-
-
-
-
-
-
-def generate_workflow_html(nodes, edges, edge_labels=None, highlight_nodes=None, highlight_edges=None, height=400):
- dot = Digraph(format='png')
-
- if edge_labels is None:
- edge_labels = {}
- if highlight_nodes is None:
- highlight_nodes = []
- if highlight_edges is None:
- highlight_edges = []
-
- node_style = {
- 'style': 'filled',
- 'fillcolor': '#DAE8FC',
- 'color': '#6C8EBF',
- 'shape': 'box',
- 'fontname': 'Helvetica',
- 'fontsize': '8',
- 'fontcolor': '#2E4057',
- 'margin': '0.1,0.02',
- 'height': '0.3',
- 'width': '0.5'
- }
-
- highlighted_node_style = {
- 'fillcolor': '#FFD580',
- 'color': '#FFA500'
- }
-
- edge_style = {
- 'color': '#6C8EBF',
- 'arrowsize': '0.5',
- 'penwidth': '0.6',
- 'fontname': 'Helvetica',
- 'fontsize': '8',
- 'fontcolor': '#2E4057',
- }
-
- highlighted_edge_style = {
- 'color': '#FFA500',
- 'penwidth': '2.0'
- }
-
- for node in nodes:
- if node in highlight_nodes:
- dot.node(node, node, **{**node_style, **highlighted_node_style})
- else:
- dot.node(node, node, **node_style)
-
- for tail, heads in edges.items():
- for head in heads:
- label = edge_labels.get((tail, head), '')
- if (tail, head) in highlight_edges:
- dot.edge(nodes[tail], nodes[head], label=label, **{**edge_style, **highlighted_edge_style})
- else:
- dot.edge(nodes[tail], nodes[head], label=label, **edge_style)
-
- png_image = dot.pipe()
-
- encoded_image = base64.b64encode(png_image).decode()
- return encoded_image
-
-def display_workflow(nodes, edges, edge_labels=None, highlight_nodes=None, highlight_edges=None, height=400):
- encoded_image = generate_workflow_html(nodes, edges, edge_labels, highlight_nodes, highlight_edges, height)
- scrollable_html = f"""
- 
- 
"""
-
- html += "
' \
- f'