Compare commits

...

No commits in common. "master" and "next" have entirely different histories.
master ... next

63 changed files with 1995 additions and 242001 deletions

1
.envrc
View file

@ -1 +0,0 @@
use flake .#impure

View file

@ -1,40 +0,0 @@
**3\. Caveats and Future Work**
Our project has some important limitations, both theoretical and empirical. The caveats below offer room for improvement in future iterations of the project, as well as new lines of investigation.
**3.1 Limitations of the theoretical approach**
**3.1.1 Temporal coherence as a concept**
Our theoretical framework is based on the concept of temporal coherence, which we characterise as the ability of putting together, consistently, different atomistic tasks in the pursuit of a more complex task. However, this is not the only possible framework for understanding AIs capabilities. For instance, some people have used the [t-AGI](https://www.alignmentforum.org/posts/BoA3agdkAzL6HQtQP/clarifying-and-predicting-agi) framework, which bears some resemblance to ours but is not identical. Also, we dont offer a formal definition of the concept. One way in which we plan to expand the project further is by incorporating it more explicitly into the model proposed by Korinek and Suh
One objection that could be raised to the concept of temporal coherence comes from the jagged-frontier of AI systems. This is the observation that AI capabilities are very irregular while they can perform tasks that would take a human many hours in mere seconds or minutes, they struggle with seemingly easy, short-term tasks. We plan to incorporate this concern in our framework in the future.
**3.1.2 Temporal coherence as a significant bottleneck**
We make an additional claim, if AIs could act coherently over arbitrary time periods of time, putting together the different atomistic tasks they are already capable of performing, more economic value would be unlocked than if any other ability was solved.
However, it is plausible to believe other bottlenecks are equally or more important. Other important bottlenecks include memory, context, multimodality and cooperation. We aim to dedicate more time to identifying the most important bottlenecks to job automation in our future research, possibly with the help of an economic model like the O-Ring model.
**3.1.3 Task length as a proxy for coherence**
Our measure for temporal coherence is the length it would take a human to complete a given task. The idea behind the proxy is that the longer the task, the more it requires putting together different pieces or atomistic tasks to complete it. But this is not a perfect measure. One way in which the proxy could be misleading is if a task takes a lot of time for a human to complete, not because it requires more coherence for putting together different subtasks, but rather because the atomistic tasks *themselves* take a lot of time to complete. To some extent, this seems contingent on how narrowly we define atomistic tasks; if done narrowly enough, as Korinek and Suh do in their original paper, then task length seems like a good proxy, because no single task would take a lot of time to complete, by definition. But theres room for disagreement, and the measure could be improved.
Finally, the concern of the jagged-frontier also applies here, because if AI systems are very irregular, its not clear that measuring the time it would take a *human* to complete a task is very informative of the temporal coherence required by *AIs.*
**3.2 Limitations of the estimates of task length**
**3.2.1 LLM classification limitations**
Humans have a [hard time](https://en.wikipedia.org/wiki/Planning_fallacy) estimating how long something takes, and LLMs too. In our test bench, we evaluated LLMs self-consistency, and found that the more the task statement seemed ambiguous, the more the model changed its estimate between a few values, but stayed consistent between those values.
We also manually estimated the time to completion for 45 task statements, and compared our results to the model, its estimates were close enough to ours.
In the future, we might try to cross-validate the estimates by developing additional methods for time estimation. There is [prior work](https://www.ons.gov.uk/economy/environmentalaccounts/articles/developingamethodformeasuringtimespentongreentasks/march2022) in this area. We might also use a model harnessed with estimation tools, or spend time developing a larger validation set so that we can evaluate our prompt and model choice more accurately.
**3.2.3 METRs estimates**
We use METRs estimates on the increasing capacity of AI agents to perform longer time-horizon tasks, doubling every 7 months, to achieve the result that more than 60 per cent of tasks can be performed by AI agents by 2026\.
The problem with using those estimates is that they are based on a set of coding benchmarks, which, as [others have commented](https://epoch.ai/gradient-updates/where-is-my-ten-minute-agi), do not reflect the messiness of real-world work environments. Also, the extent to which we can extrapolate the performance of AI agents in coding tasks to other types of tasks is an open question, especially considering how reasoning models are particularly good at coding and math, as opposed to other fields. Finally, we dont take into account the success rate variable in METRs results, which for obvious reasons seems important when it comes to automating real-world tasks.

View file

@ -1,53 +0,0 @@
**Temporal Coherence: A Bottleneck in Automation**
*This research was conducted as part of the Economics of Transformative AI hackathon with Apart Research.*
Why hasn't AI automated away more professions? One potential explanation is that, despite their intelligence, these systems cannot (yet) act coherently over long periods of time. While AI has demonstrated impressive capabilities in solving concrete, isolated problems, maintaining consistent goals, reasoning, and plans over extended time frames has proven a more significant challenge.
In this context, our project seeks to answer two questions. First, we know AI systems are getting better at completing tasks over longer time horizons, as [METRs recent work](http://.org/blog/2025-03-19-measuring-ai-ability-to-complete-long-tasks/) shows. But how much better will they need to get in order to start having a real impact on the economy, or at least for their automation capabilities to increase significantly?[^1] Second, supposing this ability is a crucial bottleneck preventing AI systems from automating economic tasks, how much value could be unlocked and how soon?
TIn order to answer these questions, we define a concept that tries to capture AIs ability to maintain consistent goals and plans while pursuing some task, called *temporal coherence*, in the context of [recent work](https://www.nber.org/papers/w32255) by Anton Korinek and Donghyun Suh. Then, we argue that temporal coherence might be the AI ability that, if solved, could unlock more economic value. With this framework in mind, we try to estimate the importance of temporal coherence in the US economy by measuring the time needed to complete all remotable tasks in the [O\*Net dataset](https://www.onetonline.org/). Our findings showsuggest that AI agents could have enough temporal coherence to perform any remote economic task by 2030, and that AI automation potential could increase in a pretty discontinuous manner, unlocking lots of economic value soon.
**What is Temporal Coherence?**
Recent discussions about the economics of AI have highlighted that while LLMs are very useful assistants and have most likely increased the productivity of many workers across the board, their effects in terms of full-job task automation have still to be felt (see [here](https://epoch.ai/epoch-after-hours/disagreements-on-agi-timelines) and [here](https://www.dwarkesh.com/p/ege-tamay)). The explanations offered for solving this puzzle usually invoke the words agency, autonomy, coherence over long horizons, or adapt plans to simple circumstances. TBecause this myriad of concepts is never clearly defined, this can create confusion.[^2]
To impose some conceptual clarity, we define the concept of *temporal coherence*. Temporal coherencehis concept is to be understood in the context of modern labor economics models, in particular, the task-based framework created by David Autor, Daron Acemoglu, Pascual Restrepo and others. In these models, tasks, rather than occupations, are the fundamental units of economic production.[^3] [Recent work](https://www.nber.org/papers/w32255) by Anton Korinek and Donghyun Suh refines this approach by modelling human work as composed of atomistic tasks that vary in computational complexity, conceptualising technological progress as expanding the automation frontier, gradually enabling machines to perform increasingly complex atomistic tasks. Our proposal is to add, on top of this atomistic framework, the ability of *putting together different atomistic tasks in the pursuit of a more complex task*. We call this ability temporal coherence.
To illustrate these abstract concepts and the usefulness of temporal coherence, let us consider teaching an economics course. While occupational databases like O\*NET might list teach economic theories as a task, this high-level label can be decomposed into subtasks: planning a syllabus, preparing lectures, delivering explanations, answering questions, recognising confusion, adjusting content dynamically, and grading assignments. Each subtask differs in computational complexity. Planning a lecture may be relatively simple, and current AI systems might already be able to do a decent job at it; dynamically adapting explanations in response to subtle student cues, however, might be more challenging.
But even if an AI system could perform each subtask—preparing slides or answering factual questions—it is very likely that current systems could not deliver a coherent semester-long course. Teaching an economics course requires doing each one of the subtasks consistently and coherently until the course is effectively completed. For instance, effective teaching requires maintaining thematic and conceptual consistency over time, adapting to cumulative student understanding, and revising instructional strategies based on longitudinal feedback. In short, AIs do not yet have enough temporal coherence to complete this task at present.
The fact that temporal coherence remains a challenge when a task involves many different atomistic tasks being put together, seems a good explanation for why we havent seen more AI task automation yet.
**The economic value of temporal coherence**
Having established the concept of temporal coherence, we now want to argue that, out of all the abilities which AI currently struggles with, solving temporal coherence would unlock the *most* economic value.
This argument comes from the observation that current AI systems are really smart, combined with an intuition about what most real-world economic activity involves in practice.
We say this because we observeThe argument for this is a combination of thethinking this is the case combines the observation that current AI systems are really smart, and we have anwith an intuition about what most real-world economic tasks involve in practice. The intelligence of state-of-the-art LLMs arecan be clearly established by the results they obtain in various benchmarks. For example, reasoning models, like [OpenAIs o3](https://openai.com/es-ES/index/introducing-o3-and-o4-mini/), seem very competent in domains like math or coding, where they get results approaching, matching or even surpassing experts in some cases.[^4] Models getting smarter and smarter rapidly is a [trend](https://ourworldindata.org/grapher/test-scores-ai-capabilities-relative-human-performance) that began with ChatGPT and continues to this day.
Despite these impressive results which clearly demonstrate AIs ability and intelligence, it does not seem that most economically valuable tasks could be performed *just* with this raw intelligence. 'Unhobblings' are necessary, to borrow Leopold Aschenbrenners [terminology](https://situational-awareness.ai/). What are AI systems lacking, exactly? Some possibilities include full multimodality, memory, context, cooperation abilities, and temporal coherence.[^5]
IConsidering just how intelligent these systems are, or, using our terminology, how many atomistic tasks they can perform, it would seem to us that if AI systems were capable of putting to use those abilities in a coherent way in the pursuit of a goal or task (that is, if they had perfect temporal coherence), then their automation potential would vastly increase. In contrast, systems with, say, perfect multimodality but without temporal coherence would suffer from the same automation limitations that LLMs face now. The automation potential of temporal coherence just seems much greater than for any other ability.[^6]
**Our Research Approach**
Temporal coherence is an ability that comes in degrees, and [recent work](http://.org/blog/2025-03-19-measuring-ai-ability-to-complete-long-tasks/) by METR shows that AI systems are getting better at it,[^7] and at a rapid pace. But this, by itself, does not tell us how important these advances are in terms of AI task automation potential. It could be that, as METRs projections show, by 2026 we have systems capable of completing tasks over an 8-hour time horizon. But if most tasks in the economy require much more coherence than this, then by 2026 temporal coherence will still be a bottleneck. To understand how big of a challenge temporal coherence will be for task automation, *we need a measure of the importance of temporal coherence in the economy*.
Unfortunately, measuring temporal coherence directly is challenginga challenge, since \--no established metric exists yet that categorises long-term projects by saying this task requires X months of coherence or that task requires Y years of coherence. For this reason, we use as a proxy for temporal coherence *the time it would take a human to complete a given task autonomously*. This is an imperfect measurement, but its supported by the simple heuristic that tasks requiring longer periods of sustained, autonomous human effort seemingly necessitate a higher degree of temporal coherence for an AI agent to successfully automate them. Further, by using this proxy, our results can be combined with METRs results to produce additional insights on the possible economic impact of future AI agents, as we explain in the next section.
[^1]: We make this clarification because many other factors are involved in actual automation of tasks. For instance, there might be other capabilities that current AI models lack which would also be needed for automation; or maybe its just not profitable to do so; or theres regulation that prohibits it; etc.
[^2]: Just as an example, in the podcast linked above from Epoch AI, Ege Erdil makes a distinction between some of these concepts that is not usually made: So to me it looks like the lack of common sense and lack of agency and ability to execute plans is a different competence than maintaining coherence over long context. This suggests clearer definitions are needed.
[^3]: See [here](https://www.nber.org/papers/w30074) for a great overview of the evolution of models about wage inequality and automation. The main papers in this area can be found [here](https://economics.mit.edu/sites/default/files/publications/the%20task%20approach%202013.pdf), [here](https://www.aeaweb.org/articles?id=10.1257/jep.33.2.3), and [here](https://economics.mit.edu/sites/default/files/publications/The%20Race%20Between%20Man%20and%20Machine%20-%20Implications%20of.pdf).
[^4]: We dont want to go over all benchmarks in this post, but some important examples include MathFrontier, SWE-bench or GPQA Diamond, in which o3 obtained impressive results according to many commentators.
[^5]: This is not meant to be an exhaustive list, in part because, again, definitions get tricky and some of these abilities might be grouped together or involve other abilities depending on how you think about them.
[^6]: Some definitional issues will be commented on Section 3\.
[^7]: Thats the way we read the evidence, at least, from the framework of temporal coherence we have proposed.

View file

@ -1,40 +0,0 @@
**Methodology**
To estimate the importance of temporal coherence in the economy, we used the O\*NET database, which lists all occupations and tasks involved infor each occupation for the US economy. We used [Barnett 2025](https://epoch.ai/gradient-updates/consequences-of-automating-remote-work)s classification of the remotable tasks in this database[^1]. The reason for this is that the automation potential from AI, at the moment, comes mostly from AI systems that can act in digital environments, rather than in the physical world.
We use a large language model to classify all remotable tasks by how much time it would take a human to perform them, which gives us a lower- and upper-bound for each task if such estimate is possible.
We exclude non-estimable tasks from the estimate, as arguably, these tasks are the ones most likely to stay human.
**![][image1]**
**Results**
Our main findings arecan be summarised byin two simple graphs. Figure 1 shows the distribution of tasks by length as classified by the LLM. The result is clear: more than 80% of tasks can be performed by humans in less than 8 hours of autonomous work. Comparatively few tasks take more than a week, and just a handful of them go over 6 months.
![][image2]
The implications of these estimates, if we take them at face value, are striking: it just doesnt seem that AIs will have to improve that much at temporal coherence for this automation bottleneck to be solved for most tasks.
This is precisely what the second graph shows. Figure 2 combines our estimates for task length with METRs projections for the increase in coherence of AI models. The result is that by 2026, less than 40% of tasks *wont* be automatable, if they depended only on temporal coherence. By the end of the decade, AI systems will have enough temporal coherence to essentially take over the US economy again, if temporal coherence were the only capability left standing. Perhaps even more importantly, the change in AI automation potential looks discrete rather than continuous, going from nearly 0% of tasks in 2025 to more than 60% in 2026\. This is due to most tasks taking around 8 hours to complete. Though we dont explore them, its clear this has important policy implications.
![][image3]
To understand how much economic value could be unlocked if temporal coherence were solved, conditional on it being the fundamental bottleneck, we investigated how much temporal coherence the five most remotable occupations required.
![][image4]
As can be seen in the graph, all top five occupations, except for Architecture and Engineering, have 80% of tasks which can be completed in a day (8 hours of work). Because AI agents could achieve this level of temporal coherence in 2026 according to METRs projections , AI agents could be unlocking a tremendous amount of economic value. To get a rough sense for the magnitude, suppose temporal coherence would be enough for task automation. Then, if 80% of tasks in all of these occupations can be automated, this could amount to *3.72$ trillion* dollars of economic value. This is of course a simplification and an upper bound, but it illustrates how quickly AI could transform the economy thanks to improvements in temporal coherence.
One question remains: how confident are we in estimates of task duration?
Half of all tasks are estimated to take between two to twenty hours, and 75% are less than or equal to 1 day on the lower end and less than or equal to 6 days on the higher end. These estimates sound reasonable for many O\*NET tasks.
We find that 90% of task duration estimates have an upper to lower bound duration ratio lower than or equal to 10\.
![][image5]
Tasks with very large ratios accurately reflect the ambiguity of the task description. For instance, the task Conduct research in a particular field of knowledge and publish findings in professional journals, books, or electronic media takes between 40 hours and 1 year (216x difference). Another example is the task Create or use statistical models for the analysis of genetic data, which takes between one day and six months (180x difference).
**The value of our results**
Why should you care about these results?. First, our estimates of the distribution of task length for the US economy point to temporal coherence or whatever name you want to give to the capacity of people to work autonomously on a task until completion being a *less important factor* than we expected. The fact that 80% of tasks can be completed in 8 hours or less suggests that, regardless of the actual importance of temporal coherence for unlocking economic value, this will not be a bottleneck for very long. When we combine our estimates with METRs, which project the increasing capacity of AI agents to perform tasks over longer time horizons, we realise that by 2026, temporal coherence wont be an issue for automating more than 60% of all tasks.
Second, once you add in the assumption that temporal coherence is likely to be the biggest bottleneck of all the abilities that AIs currently lack (or have not perfected): Our results suggest that a lot of economic value could potentially be unlocked due to increases in AI agents temporal coherence,[^2] and that this could happen fairly soon. This has potential policy implications that we do not explore in this piece, but seem particularly relevant for concerns about [gradual disempowerment.](https://gradual-disempowerment.ai/)
[^1]: We think Barnetts classification could be improved further; for example, by using O\*NETs [Physical Work Conditions](https://www.onetcenter.org/content.html#cm-sub-4-C-2) annotations.
[^2]: This is not guaranteed, however, because actual automation depends on other factors which we dont analyze here, like the cost of running the AI agents and integrating them into current business, regulation and other frictions.

View file

@ -1,507 +0,0 @@
import pandas as pd
import litellm
import dotenv
import os
import time
import json
import math
import numpy as np
# --- Configuration ---
MODEL = "gpt-4.1-mini" # Make sure this model supports json_schema or structured output
RATE_LIMIT = 5000 # Requests per minute
CHUNK_SIZE = 300
SECONDS_PER_MINUTE = 60
FILENAME = (
"tasks_with_estimates.csv" # This CSV should contain the tasks to be processed
)
# --- Prompts and Schema ---
SYSTEM_PROMPT = """
You are an expert assistant evaluating the time to completion required for job tasks. Your goal is to estimate the time range needed for a skilled human to complete the following job task remotely, without supervision.
Provide a lower and upper bound estimate for the time to completion time. These bounds should capture the time within which approximately 80% of instances of performing this specific task are typically completed by a qualified individual.
Base your estimate on the provided task description, its associated activities, and the occupational context. Your estimate must be in one the allowed units: minute, hour, day, week, month, trimester, semester, year.
""".strip()
USER_MESSAGE_TEMPLATE = """
Please estimate the time range for the following remote task:
**Task Description:** {task}
**Relevant activies for the task:**
{dwas}
**Occupation Category:** {occupation_title}
**Occupation Description:** {occupation_description}
Consider the complexity and the typical steps involved.
""".strip()
ALLOWED_UNITS = [
"minute",
"hour",
"day",
"week",
"month",
"trimester",
"semester",
"year",
]
SCHEMA_FOR_VALIDATION = {
"name": "estimate_time",
"strict": True, # Enforce schema adherence
"schema": {
"type": "object",
"properties": {
"lower_bound_estimate": {
"type": "object",
"properties": {
"quantity": {
"type": "number",
"description": "The numerical value for the lower bound of the estimate.",
},
"unit": {
"type": "string",
"enum": ALLOWED_UNITS,
"description": "The unit of time for the lower bound.",
},
},
"required": ["quantity", "unit"],
"additionalProperties": False,
},
"upper_bound_estimate": {
"type": "object",
"properties": {
"quantity": {
"type": "number",
"description": "The numerical value for the upper bound of the estimate.",
},
"unit": {
"type": "string",
"enum": ALLOWED_UNITS,
"description": "The unit of time for the upper bound.",
},
},
"required": ["quantity", "unit"],
"additionalProperties": False,
},
},
"required": ["lower_bound_estimate", "upper_bound_estimate"],
"additionalProperties": False,
},
}
def save_dataframe(df_to_save, filename):
"""Saves the DataFrame to the specified CSV file using atomic write."""
try:
temp_filename = filename + ".tmp"
df_to_save.to_csv(temp_filename, encoding="utf-8-sig", index=False)
os.replace(temp_filename, filename)
except Exception as e:
print(f"--- Error saving DataFrame to {filename}: {e} ---")
if os.path.exists(temp_filename):
try:
os.remove(temp_filename)
except Exception as remove_err:
print(
f"--- Error removing temporary save file {temp_filename}: {remove_err} ---"
)
def create_task_estimates():
try:
# Read the CSV
if os.path.exists(FILENAME):
df = pd.read_csv(FILENAME, encoding="utf-8-sig")
print(f"Successfully read {len(df)} rows from {FILENAME}.")
estimate_columns_spec = {
"lb_estimate_qty": float,
"lb_estimate_unit": object,
"ub_estimate_qty": float,
"ub_estimate_unit": object,
}
save_needed = False
for col_name, target_dtype in estimate_columns_spec.items():
if col_name not in df.columns:
# Initialize with a type-compatible missing value
if target_dtype == float:
df[col_name] = np.nan
else: # object
df[col_name] = pd.NA
df[col_name] = df[col_name].astype(target_dtype) # Enforce dtype
print(f"Added '{col_name}' column as {df[col_name].dtype}.")
save_needed = True
else:
# Column exists, ensure correct dtype
current_pd_dtype = df[col_name].dtype
expected_pd_dtype = pd.Series(dtype=target_dtype).dtype
if current_pd_dtype != expected_pd_dtype:
try:
if target_dtype == float:
df[col_name] = pd.to_numeric(df[col_name], errors="coerce")
else: # object
df[col_name] = df[col_name].astype(object)
print(
f"Corrected dtype of '{col_name}' to {df[col_name].dtype}."
)
save_needed = True
except Exception as e:
print(
f"Warning: Could not convert column '{col_name}' to {target_dtype}: {e}. Current dtype: {current_pd_dtype}"
)
# Standardize missing values (e.g., empty strings to NA/NaN)
# Replace common missing placeholders with pd.NA first
df[col_name].replace(["", None, ""], pd.NA, inplace=True)
if target_dtype == float:
# For float columns, ensure they are numeric and use np.nan after replacement
df[col_name] = pd.to_numeric(df[col_name], errors="coerce")
if save_needed:
print(f"Saving {FILENAME} after adding/adjusting estimate columns.")
save_dataframe(df, FILENAME)
else:
print(
f"Error: {FILENAME} not found. Please ensure the file exists and contains task data."
)
exit()
except FileNotFoundError:
print(
f"Error: {FILENAME} not found. Please ensure the file exists and contains task data."
)
exit()
except Exception as e:
print(f"Error reading or initializing {FILENAME}: {e}")
exit()
# --- Identify Rows to Process ---
# We'll check for NaN in one of the primary quantity columns.
unprocessed_mask = df["lb_estimate_qty"].isna()
if unprocessed_mask.any():
start_index = unprocessed_mask.idxmax() # Finds the index of the first True value
print(f"Resuming processing. First unprocessed row found at index {start_index}.")
df_to_process = df.loc[unprocessed_mask].copy()
original_indices = df_to_process.index # Keep track of original indices
else:
print(
"All rows seem to have estimates already (based on 'lb_estimate_qty'). Exiting."
)
exit()
# --- Prepare messages for batch completion (only for rows needing processing) ---
messages_list = []
skipped_rows_indices = []
valid_original_indices = []
if not df_to_process.empty:
required_cols = ["task", "occupation_title", "occupation_description", "dwas"]
print(
f"Preparing messages for up to {len(df_to_process)} rows starting from original index {original_indices[0] if len(original_indices) > 0 else 'N/A'}..."
)
print(f"Checking for required columns: {required_cols}")
for index, row in df_to_process.iterrows():
missing_or_empty = []
for col in required_cols:
if col not in row or pd.isna(row[col]) or str(row[col]).strip() == "":
missing_or_empty.append(col)
if missing_or_empty:
print(
f"Warning: Skipping row original index {index} due to missing/empty required data in columns: {', '.join(missing_or_empty)}."
)
skipped_rows_indices.append(index)
continue
try:
user_message = USER_MESSAGE_TEMPLATE.format(
task=row["task"],
occupation_title=row["occupation_title"],
occupation_description=row["occupation_description"],
dwas=row["dwas"],
)
except KeyError as e:
print(
f"Error: Skipping row original index {index} due to formatting error - missing key: {e}. Check USER_MESSAGE_TEMPLATE and CSV columns."
)
skipped_rows_indices.append(index)
continue
messages_for_row = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_message},
]
messages_list.append(messages_for_row)
valid_original_indices.append(index) # This is the original DataFrame index
print(
f"Prepared {len(messages_list)} valid message sets for batch completion (skipped {len(skipped_rows_indices)} rows)."
)
if not messages_list:
print("No valid rows found to process after checking required data. Exiting.")
exit()
else:
print(
"No rows found needing processing (df_to_process is empty)."
) # Should have been caught by earlier check
exit()
# --- Call batch_completion in chunks with rate limiting and periodic saving ---
total_messages_to_send = len(messages_list)
num_chunks = math.ceil(total_messages_to_send / CHUNK_SIZE)
print(
f"\nStarting batch completion for {total_messages_to_send} items in {num_chunks} chunks..."
)
overall_start_time = time.time()
processed_count_total = 0
for i in range(num_chunks):
chunk_start_message_index = i * CHUNK_SIZE
chunk_end_message_index = min((i + 1) * CHUNK_SIZE, total_messages_to_send)
message_chunk = messages_list[chunk_start_message_index:chunk_end_message_index]
# Get corresponding original DataFrame indices for this chunk
chunk_original_indices = valid_original_indices[
chunk_start_message_index:chunk_end_message_index
]
if not message_chunk:
continue
min_idx_disp = min(chunk_original_indices) if chunk_original_indices else "N/A"
max_idx_disp = max(chunk_original_indices) if chunk_original_indices else "N/A"
print(
f"\nProcessing chunk {i + 1}/{num_chunks} (Messages {chunk_start_message_index + 1}-{chunk_end_message_index} of this run)..."
f" Corresponding to original indices: {min_idx_disp} - {max_idx_disp}"
)
chunk_start_time = time.time()
responses = []
try:
print(f"Sending {len(message_chunk)} requests for chunk {i + 1}...")
responses = litellm.batch_completion(
model=MODEL,
messages=message_chunk,
response_format={
"type": "json_schema",
"json_schema": SCHEMA_FOR_VALIDATION,
},
num_retries=3,
# request_timeout=60 # Optional: uncomment if needed
)
print(f"Chunk {i + 1} API call completed.")
except Exception as e:
print(f"Error during litellm.batch_completion for chunk {i + 1}: {e}")
responses = [None] * len(
message_chunk
) # Ensure responses list matches message_chunk length for processing loop
# --- Process responses for the current chunk ---
chunk_updates = {} # To store {original_df_index: {qty/unit data}}
successful_in_chunk = 0
failed_in_chunk = 0
if responses and len(responses) == len(message_chunk):
for j, response in enumerate(responses):
original_df_index = chunk_original_indices[j]
# Initialize values for this item
lb_qty_val, lb_unit_val, ub_qty_val, ub_unit_val = None, None, None, None
content_str = None
if response is None:
print(
f"Skipping processing for original index {original_df_index} due to API call failure for this item (response is None)."
)
failed_in_chunk += 1
continue
try:
if (
response.choices
and response.choices[0].message
and response.choices[0].message.content
):
content_str = response.choices[0].message.content
estimate_data = json.loads(content_str) # Can raise JSONDecodeError
lower_bound_dict = estimate_data.get("lower_bound_estimate")
upper_bound_dict = estimate_data.get("upper_bound_estimate")
valid_response_structure = isinstance(
lower_bound_dict, dict
) and isinstance(upper_bound_dict, dict)
if valid_response_structure:
lb_qty_raw = lower_bound_dict.get("quantity")
lb_unit_raw = lower_bound_dict.get("unit")
ub_qty_raw = upper_bound_dict.get("quantity")
ub_unit_raw = upper_bound_dict.get("unit")
is_valid_item = True
# Validate LB Qty
if (
not isinstance(lb_qty_raw, (int, float))
or math.isnan(float(lb_qty_raw))
or float(lb_qty_raw) < 0
):
print(
f"Warning: Invalid lb_quantity for original index {original_df_index}: {lb_qty_raw}"
)
is_valid_item = False
else:
lb_qty_val = float(lb_qty_raw)
# Validate UB Qty
if (
not isinstance(ub_qty_raw, (int, float))
or math.isnan(float(ub_qty_raw))
or float(ub_qty_raw) < 0
):
print(
f"Warning: Invalid ub_quantity for original index {original_df_index}: {ub_qty_raw}"
)
is_valid_item = False
else:
ub_qty_val = float(ub_qty_raw)
# Validate Units
if lb_unit_raw not in ALLOWED_UNITS:
print(
f"Warning: Invalid lb_unit for original index {original_df_index}: '{lb_unit_raw}'"
)
is_valid_item = False
else:
lb_unit_val = lb_unit_raw
if ub_unit_raw not in ALLOWED_UNITS:
print(
f"Warning: Invalid ub_unit for original index {original_df_index}: '{ub_unit_raw}'"
)
is_valid_item = False
else:
ub_unit_val = ub_unit_raw
if is_valid_item:
successful_in_chunk += 1
chunk_updates[original_df_index] = {
"lb_estimate_qty": lb_qty_val,
"lb_estimate_unit": lb_unit_val,
"ub_estimate_qty": ub_qty_val,
"ub_estimate_unit": ub_unit_val,
}
else:
failed_in_chunk += (
1 # Values remain None if not fully valid
)
else:
print(
f"Warning: Missing or malformed estimate dicts in JSON for original index {original_df_index}. Content: '{content_str}'"
)
failed_in_chunk += 1
else:
finish_reason = (
response.choices[0].finish_reason
if (response.choices and response.choices[0].finish_reason)
else "unknown"
)
error_message = (
response.choices[0].message.content
if (
response.choices
and response.choices[0].message
and response.choices[0].message.content
)
else "No content in message."
)
print(
f"Warning: Received non-standard or empty response content for original index {original_df_index}. "
f"Finish Reason: '{finish_reason}'. Message: '{error_message}'. Raw Choices: {response.choices}"
)
failed_in_chunk += 1
except json.JSONDecodeError:
print(
f"Warning: Could not decode JSON for original index {original_df_index}. Content received: '{content_str}'"
)
failed_in_chunk += 1
except AttributeError as ae:
print(
f"Warning: Missing expected attribute processing response for original index {original_df_index}: {ae}. Response: {response}"
)
failed_in_chunk += 1
except Exception as e:
print(
f"Warning: An unexpected error occurred processing response for original index {original_df_index}: {type(e).__name__} - {e}. Response: {response}"
)
failed_in_chunk += 1
else:
print(
f"Warning: Mismatch between number of responses ({len(responses) if responses else 0}) "
f"and messages sent ({len(message_chunk)}) for chunk {i + 1}, or no responses. Marking all as failed."
)
failed_in_chunk = len(
message_chunk
) # All items in this chunk are considered failed if response array is problematic
print(
f"Chunk {i + 1} processing summary: Success={successful_in_chunk}, Failed/Skipped={failed_in_chunk}"
)
processed_count_total += successful_in_chunk
# --- Update Main DataFrame and Save Periodically ---
if chunk_updates:
print(
f"Updating main DataFrame with {len(chunk_updates)} new estimates for chunk {i + 1}..."
)
for idx, estimates in chunk_updates.items():
if idx in df.index:
df.loc[idx, "lb_estimate_qty"] = estimates["lb_estimate_qty"]
df.loc[idx, "lb_estimate_unit"] = estimates["lb_estimate_unit"]
df.loc[idx, "ub_estimate_qty"] = estimates["ub_estimate_qty"]
df.loc[idx, "ub_estimate_unit"] = estimates["ub_estimate_unit"]
print(f"Saving progress to {FILENAME}...")
save_dataframe(df, FILENAME)
else:
print(f"No successful estimates obtained in chunk {i + 1} to save.")
# --- Rate Limiting Pause ---
chunk_end_time = time.time()
chunk_duration = chunk_end_time - chunk_start_time
print(f"Chunk {i + 1} took {chunk_duration:.2f} seconds.")
if i < num_chunks - 1: # No pause after the last chunk
# Calculate ideal time per request based on rate limit
time_per_request = SECONDS_PER_MINUTE / RATE_LIMIT if RATE_LIMIT > 0 else 0
# Calculate minimum duration this chunk should have taken to respect rate limit
min_chunk_duration_for_rate = len(message_chunk) * time_per_request
# Calculate pause needed
pause_needed = max(0, min_chunk_duration_for_rate - chunk_duration)
if pause_needed > 0:
print(
f"Pausing for {pause_needed:.2f} seconds to respect rate limit ({RATE_LIMIT}/min)..."
)
time.sleep(pause_needed)
overall_end_time = time.time()
total_duration_minutes = (overall_end_time - overall_start_time) / 60
print(
f"\nBatch completion finished."
f" Processed {processed_count_total} new estimates in this run in {total_duration_minutes:.2f} minutes."
)
print(f"Performing final save to {FILENAME}...")
save_dataframe(df, FILENAME)
print("\nScript finished.")

View file

@ -1,2 +1,3 @@
- I use Nix. To run a command, prefix them with `nix develop .#impure -c`
- I use uv. To add a package, use: uv add. To run a script use: uv run path/to/script
- To run the pipeline: `uv run -m pipeline.runner`

View file

@ -1,563 +0,0 @@
import os
import litellm
import sqlite3
import numpy as np
import pandas as pd
from google.colab import userdata, files
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib as mpl
os.environ['OPENAI_API_KEY'] = userdata.get('OPENAI_API_KEY')
os.environ['GEMINI_API_KEY'] = userdata.get('GEMINI_API_KEY')
occupation_major_codes = {
'11': 'Management',
'13': 'Business and Financial Operations',
'15': 'Computer and Mathematical Occupations',
'17': 'Architecture and Engineering',
'19': 'Life, Physical, and Social Science',
'21': 'Community and Social Services',
'23': 'Legal',
'25': 'Education, Training, and Library',
'27': 'Arts, Design, Entertainment, Sports, and Media',
'29': 'Healthcare Practitioners and Technical',
'31': 'Healthcare Support',
'33': 'Protective Service',
'35': 'Food Preparation and Serving Related',
'37': 'Building and Grounds Cleaning and Maintenance',
'39': 'Personal Care and Service',
'41': 'Sales and Related',
'43': 'Office and Administrative Support',
'45': 'Farming, Fishing, and Forestry',
'47': 'Construction and Extraction',
'49': 'Installation, Maintenance, and Repair',
'51': 'Production',
'53': 'Transportation and Material Moving',
'55': 'Military Specific'
}
gray = {'50':'#f8fafc','100':'#f1f5f9','200':'#e2e8f0',
'300':'#cbd5e1','400':'#94a3b8','500':'#64748b',
'600':'#475569','700':'#334155','800':'#1e293b',
'900':'#0f172a','950':'#020617'}
lime = {'50': '#f7fee7','100': '#ecfcca','200': '#d8f999',
'300': '#bbf451','400': '#9ae600','500': '#83cd00',
'600': '#64a400','700': '#497d00','800': '#3c6300',
'900': '#35530e','950': '#192e03'}
mpl.rcParams.update({
'figure.facecolor' : gray['50'],
'axes.facecolor' : gray['50'],
'axes.edgecolor' : gray['100'],
'axes.labelcolor' : gray['700'],
'xtick.color' : gray['700'],
'ytick.color' : gray['700'],
'font.family' : 'Inter', # falls back to DejaVu if Inter not present
'font.size' : 11,
})
sns.set_style("white") # keep minimal axes, we will remove default grid
sns.set_context("notebook")
def prepare_tasks():
# This dataset comes from https://epoch.ai/gradient-updates/consequences-of-automating-remote-work
# It contains labels for a O*NET task can be done remotely or not (labeled by GPT-4o)
# You can download it here: https://drive.google.com/file/d/1GrHhuYIgaCCgo99dZ_40BWraz-fzo76r/view?usp=sharing
df_remote_status = pd.read_csv("epoch_task_data.csv")
# BLS OEWS: https://www.bls.gov/oes/special-requests/oesm23nat.zip
df_oesm = pd.read_excel("oesm23national.xlsx")
# Run uv run ./enrich_task_ratings.py
df_tasks = pd.read_json("task_ratings_enriched.json")
# Run uv run classify_estimateability_of_tasks.py
df_task_estimateable = pd.read_csv("tasks_estimateable.csv").rename(columns={"task_estimateable": "estimateable"}).drop_duplicates(subset=['task'], keep='first')
# df_tasks now has a remote_status column which contains either "remote" or "not remote"
df_tasks = pd.merge(df_tasks, df_remote_status[['Task', 'Remote']], left_on='task', right_on='Task', how='left')
df_tasks = df_tasks.drop('Task', axis=1).rename(columns={'Remote': 'remote_status'})
# df_tasks now has a estimateable column which contains either "ATOMIC" or "ONGOING-CONSTRAINT"
df_tasks = pd.merge(df_tasks, df_task_estimateable[['task', 'estimateable']], on='task', how='left')
df_tasks = df_tasks[df_tasks['importance_average'] < 3].copy()
df_tasks['onetsoc_major'] = df_tasks['onetsoc_code'].str[:2]
df_remote_tasks = df_tasks[df_tasks['remote_status'] == 'remote'].copy()
# Call create_task_estimates() from add_task_estimates? which creates tasks_with_estimates.csv
def preprocessing_time_estimates():
df = pd.read_csv("tasks_with_estimates.csv")
df = df[df['importance_average'] > 3].copy()
# The embeddings comes from running `uv run ./embed_task_description.py`
# Columns: ['embedding_id', 'task', 'embedding_vector']
# These contain embedding for UNIQUE tasks
df_task_embeddings = pd.read_parquet("tasks_with_embeddings.parquet").drop_duplicates(subset=['task'])[['task', 'task_embedding']].rename(columns={"task_embedding": "embedding_vector"}).copy()
df = pd.merge(df, df_task_embeddings[['task', 'embedding_vector']], on='task', how='left')
df = pd.merge(df, df_task_estimateable[['task', 'estimateable']], on='task', how='left')
df['onetsoc_major'] = df['onetsoc_code'].str[:2]
def convert_to_minutes(qty, unit):
"""Converts a quantity in a given unit to minutes."""
return qty * {
"minute": 1,
"hour": 60,
"day": 60 * 24,
"week": 60 * 24 * 7,
"month": 60 * 24 * 30,
"trimester": 60 * 24 * 90,
"semester": 60 * 24 * 180,
"year": 60 * 24 * 365,
}[unit]
df['lb_estimate_in_minutes'] = df.apply(
lambda row: convert_to_minutes(row['lb_estimate_qty'], row['lb_estimate_unit']), axis=1
)
df['ub_estimate_in_minutes'] = df.apply(
lambda row: convert_to_minutes(row['ub_estimate_qty'], row['ub_estimate_unit']), axis=1
)
df['estimate_range'] = df.ub_estimate_in_minutes - df.lb_estimate_in_minutes
df['estimate_ratio'] = df.ub_estimate_in_minutes / df.lb_estimate_in_minutes
df['estimate_midpoint'] = (df.lb_estimate_in_minutes + df.ub_estimate_in_minutes)/2
atomic_tasks = df[df['estimateable'] == 'ATOMIC']
ongoing_tasks = df[df['estimateable'] == 'ONGOING-CONSTRAINT']
with pd.option_context('display.max_columns', None):
display(df)
# Check for empty estimates
if atomic_tasks['lb_estimate_in_minutes'].isnull().sum() > 0:
print("Missing values in 'lb_estimate_in_minutes':", atomic_tasks['lb_estimate_in_minutes'].isnull().sum())
if atomic_tasks['ub_estimate_in_minutes'].isnull().sum() > 0:
print("Missing values in 'ub_estimate_in_minutes':", atomic_tasks['ub_estimate_in_minutes'].isnull().sum())
# Check for impossible bounds
impossible_bounds = atomic_tasks[
(atomic_tasks['lb_estimate_in_minutes'] <= 0) |
(atomic_tasks['ub_estimate_in_minutes'] <= 0) |
(atomic_tasks['lb_estimate_in_minutes'] > atomic_tasks['ub_estimate_in_minutes'])
]
if not impossible_bounds.empty:
print(f"Error: Found rows with impossible bounds.")
with pd.option_context('display.max_colwidth', None):
display(impossible_bounds[['task', 'lb_estimate_in_minutes', 'ub_estimate_in_minutes', 'dwas']])
#with pd.option_context('display.max_colwidth', None):
#display(atomic_tasks.nlargest(20, 'ub_estimate_in_minutes')[['task', 'lb_estimate_qty', 'lb_estimate_unit', 'lb_estimate_in_minutes', 'ub_estimate_qty', 'ub_estimate_unit', 'ub_estimate_in_minutes', 'estimate_ratio']])
def cell1():
sns.histplot(atomic_tasks.estimate_midpoint, log_scale=True)
def cell2():
plt.figure(figsize=(14,10))
sns.boxplot(
data=atomic_tasks,
x='onetsoc_major', # 11 = Management, 15 = Computer/Math, …
y='estimate_range',
showfliers=False
)
plt.yscale('log') # long tail => log scale
plt.xlabel('Occupation')
plt.ylabel('Range (upper-lower, minutes)')
plt.title('Spread of time-range estimates per occupation')
ax = plt.gca()
ax.set_xticklabels([occupation_major_codes[code.get_text()] for code in ax.get_xticklabels()], rotation=60, ha='right')
def cell3():
plt.figure(figsize=(10, 10))
ax = sns.scatterplot(
data=atomic_tasks.replace({'onetsoc_major': occupation_major_codes}), # Replace codes with labels
x='lb_estimate_in_minutes', y='ub_estimate_in_minutes',
alpha=0.2, edgecolor=None, hue="onetsoc_major" # Use the labeled column for hue
)
# 45° reference
lims = (1, atomic_tasks[['lb_estimate_in_minutes','ub_estimate_in_minutes']].max().max())
ax.plot(lims, lims, color='black', linestyle='--', linewidth=1)
# optional helper lines: 2× and 10×, 100× ratios
for k in [2,10, 100]:
ax.plot(lims, [k*l for l in lims],
linestyle=':', color='grey', linewidth=1)
ax.set(xscale='log', yscale='log')
ax.set_xlabel('Lower-bound (min, log scale)')
ax.set_ylabel('Upper-bound (min, log scale)')
ax.set_title('Lower vs upper estimates for all tasks')
# Place the legend outside the plot
ax.legend(bbox_to_anchor=(1, 1), loc='upper left')
def cell4():
plt.figure(figsize=(8,4))
sns.histplot(np.log10(atomic_tasks['estimate_ratio'].replace([np.inf, -np.inf], np.nan).dropna()),
bins=60, kde=True)
plt.axvline(np.log10(10), color='red', ls='--', lw=1, label='10×')
plt.axvline(np.log10(1.05), color='orange', ls='--', lw=1, label='1.05×')
plt.axvline(0, color='black', ls='-', lw=1) # ub = lb
plt.xlabel('log₁₀(upper / lower)')
plt.ylabel('Count')
plt.title('Distribution of upper:lower ratio')
plt.legend()
plt.tight_layout()
def cell5():
# 1. Bin lower bounds into quartiles (Q1Q4)
atomic_tasks['lb_q'] = pd.qcut(atomic_tasks.lb_estimate_in_minutes,
q=4, labels=['Q1 shortest','Q2','Q3','Q4 longest'])
# 3. Aggregate: median (or mean) ratio per cell
pivot = atomic_tasks.pivot_table(index='onetsoc_major', columns='lb_q',
values='estimate_ratio', aggfunc='median')
# Map the index (onetsoc_major codes) to their corresponding labels
pivot.index = pivot.index.map(occupation_major_codes)
# 4. Visualise
plt.figure(figsize=(10,8))
sns.heatmap(pivot, cmap='RdYlGn_r', center=2, annot=True, fmt='.1f',
cbar_kws={'label':'Median upper/lower ratio'})
plt.xlabel('Lower-bound quartile')
plt.ylabel('Occupation (major group)')
plt.title('Typical range width by occupation and task length')
plt.tight_layout()
def cell6():
"""
from scipy.stats import median_abs_deviation
def mad_z(series):
med = series.median()
mad = median_abs_deviation(series, scale='normal') # ⇒ comparable to σ
return (series - med) / mad
df['robust_z'] = df.groupby('onetsoc_code')['estimate_midpoint'].transform(mad_z)
"""
agg = (atomic_tasks
.groupby('onetsoc_code')['estimate_midpoint']
.agg(median='median',
q1=lambda x: x.quantile(.25),
q3=lambda x: x.quantile(.75),
mean='mean',
std='std')
.reset_index())
agg['IQR'] = agg.q3 - agg.q1
agg['CV'] = agg['std'] / agg['mean'] # coefficient of variation
# merge back the group mean and std so each row can be scored
atomic_tasks = atomic_tasks.merge(agg[['onetsoc_code','mean','std']], on='onetsoc_code')
atomic_tasks['z'] = (atomic_tasks.estimate_midpoint - atomic_tasks['mean']) / atomic_tasks['std']
outliers = atomic_tasks.loc[atomic_tasks.z.abs() > 3]
outliers
def cell7():
from scipy.stats import median_abs_deviation
def mad_z(series):
med = series.median()
mad = median_abs_deviation(series, scale='normal') # ⇒ comparable to σ
return (series - med) / mad
atomic_tasks['robust_z'] = atomic_tasks.groupby('onetsoc_code')['estimate_midpoint'].transform(mad_z)
def cell8():
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.neighbors import NearestNeighbors
# unit-normalise embeddings
X = np.vstack(atomic_tasks.embedding_vector.values)
X = X / np.linalg.norm(X, axis=1, keepdims=True)
k = 5
nn = NearestNeighbors(n_neighbors=k+1, metric='cosine').fit(X)
dist, idx = nn.kneighbors(X) # idx[:,0] is the task itself
pairs = pd.DataFrame({
'i': np.repeat(np.arange(len(X)), k),
'j': idx[:,1:].ravel(), # drop self-match
'sim': 1 - dist[:,1:].ravel() # cosine similarity
})
# join time estimates
pairs = (pairs
.merge(atomic_tasks[['estimate_midpoint']], left_on='i', right_index=True)
.merge(atomic_tasks[['estimate_midpoint']], left_on='j', right_index=True,
suffixes=('_i','_j')))
pairs['ratio'] = pairs[['estimate_midpoint_i','estimate_midpoint_j']].max(1) / \
pairs[['estimate_midpoint_i','estimate_midpoint_j']].min(1)
pairs
def cell9():
sns.scatterplot(data=pairs.sample(20_000), x='sim', y=np.log10(pairs['ratio']),
alpha=.1)
plt.axhline(np.log10(2), ls=':', color='red'); # e.g. >2× difference
plt.xlabel('Cosine similarity'); plt.ylabel('log₁₀ time-ratio');
plt.title('Are similar tasks given similar time estimates?');
def cell10():
import matplotlib.ticker as mtick # For percentage formatting
import matplotlib.colors as mcolors # For color conversion
summary_data = []
for code, label in occupation_major_codes.items():
occ_df = df_tasks[df_tasks['onetsoc_major'] == code]
total_tasks_in_occ = len(occ_df)
if total_tasks_in_occ == 0:
continue # Skip if no tasks for this occupation
# Stack 1: % that isn't equal to "remote"
not_remote_count = len(occ_df[occ_df['remote_status'] != 'remote'])
# For the remaining remote tasks:
remote_df = occ_df[occ_df['remote_status'] == 'remote']
# Stack 2: % of remote + ATOMIC
remote_atomic_count = len(remote_df[remote_df['estimateable'] == 'ATOMIC'])
# Stack 3: % of remote + ONGOING-CONSTRAINT
remote_ongoing_count = len(remote_df[remote_df['estimateable'] == 'ONGOING-CONSTRAINT'])
summary_data.append({
'onetsoc_major_code': code,
'occupation_label': label,
'count_not_remote': not_remote_count,
'count_remote_atomic': remote_atomic_count,
'count_remote_ongoing': remote_ongoing_count,
'total_tasks': total_tasks_in_occ
})
summary_df = pd.DataFrame(summary_data)
# --- 3. Calculate Percentages ---
# Ensure total_tasks is not zero to avoid division by zero errors if an occupation had no tasks
summary_df = summary_df[summary_df['total_tasks'] > 0].copy() # Use .copy() to avoid SettingWithCopyWarning
summary_df['pct_not_remote'] = (summary_df['count_not_remote'] / summary_df['total_tasks']) * 100
summary_df['pct_remote_atomic'] = (summary_df['count_remote_atomic'] / summary_df['total_tasks']) * 100
summary_df['pct_remote_ongoing'] = (summary_df['count_remote_ongoing'] / summary_df['total_tasks']) * 100
# Select columns for plotting and set index to occupation label
plot_df = summary_df.set_index('occupation_label')[
['pct_not_remote', 'pct_remote_atomic', 'pct_remote_ongoing']
]
# Rename columns for a clearer legend
plot_df.columns = ['Not Remote', 'Remote + Estimable', 'Remote + Not estimable']
plot_df = plot_df.sort_values(by='Not Remote', ascending=False)
# --- 4. Plotting (Modified) ---
# Define the custom colors based on your requirements
# The order must match the column order in plot_df:
# 1. 'Not Remote'
# 2. 'Remote & ATOMIC'
# 3. 'Remote & ONGOING-CONSTRAINT'
bar_colors = [gray["300"], lime["500"], lime["200"]]
fig, ax = plt.subplots(figsize=(14, 10)) # Adjusted figsize for better readability
plot_df.plot(kind='barh', stacked=True, ax=ax, color=bar_colors)
ax.set_xlabel("Percentage of Tasks (%)", fontsize=12)
ax.set_ylabel("Occupation Major Group", fontsize=12)
ax.set_title("Task Breakdown by Occupation, Remote Status, and Estimateability", fontsize=14, pad=20)
# Format x-axis as percentages
ax.xaxis.set_major_formatter(mtick.PercentFormatter())
plt.xlim(0, 100) # Ensure x-axis goes from 0 to 100%
# Remove right and top spines
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
# Function to get contrasting text color
def get_contrasting_text_color(bg_color_hex_or_rgba):
"""
Determines if black or white text provides better contrast against a given background color.
bg_color_hex_or_rgba: A hex string (e.g., '#RRGGBB') or an RGBA tuple (values in [0, 1]).
Returns: 'black' or 'white'.
"""
# Convert to RGBA if it's a hex string or name
if isinstance(bg_color_hex_or_rgba, str):
rgba = mcolors.to_rgba(bg_color_hex_or_rgba)
else:
rgba = bg_color_hex_or_rgba
r, g, b, _ = rgba # Ignore alpha for luminance calculation
# Calculate luminance (standard formula for sRGB)
# Values r, g, b should be in [0, 1] for this formula
luminance = 0.2126 * r + 0.7152 * g + 0.0722 * b
# Threshold for deciding text color
return 'black' if luminance > 0.55 else 'white' # Adjusted threshold slightly for better visual
# Add percentages inside each bar segment
# Iterate through each "category" of bars (Not Remote, Remote & ATOMIC, etc.)
for i, container in enumerate(ax.containers):
# Get the color for this container/category
segment_color = bar_colors[i]
text_color = get_contrasting_text_color(segment_color)
for patch in container.patches: # Iterate through each bar segment in the category
width = patch.get_width()
if width > 3: # Only add text if segment is wide enough (e.g., >3%)
x = patch.get_x() + width / 2
y = patch.get_y() + patch.get_height() / 2
ax.text(x, y,
f"{width:.1f}%",
ha='center',
va='center',
fontsize=8, # Adjust font size as needed
color=text_color,
fontweight='medium') # Bolder text can help
plt.legend(title="Task Category", bbox_to_anchor=(1.02, 1), loc='upper left', frameon=False)
def cell11():
df_oesm['onetsoc_major'] = df_oesm['OCC_CODE'].str[:2]
# Calculate wage bill per occupation
# Wage bill = Total Employment * Annual Mean Wage
# Ensure columns are numeric, converting non-numeric values to NaN first
df_oesm['TOT_EMP'] = pd.to_numeric(df_oesm['TOT_EMP'], errors='coerce')
df_oesm['A_MEAN'] = pd.to_numeric(df_oesm['A_MEAN'], errors='coerce')
# Drop rows with NaN in necessary columns after coercion
df_oesm.dropna(subset=['TOT_EMP', 'A_MEAN', 'onetsoc_major'], inplace=True)
df_oesm['wage_bill'] = df_oesm['TOT_EMP'] * df_oesm['A_MEAN']
# Aggregate wage bill by onetsoc_major
df_wage_bill_major = df_oesm.groupby('onetsoc_major')['wage_bill'].sum().reset_index()
# Map major codes to titles for better plotting
df_wage_bill_major['OCC_TITLE_MAJOR'] = df_wage_bill_major['onetsoc_major'].map(occupation_major_codes)
# Sort by wage bill for better visualization
df_wage_bill_major = df_wage_bill_major.sort_values('wage_bill', ascending=False)
# Plotting
plt.figure(figsize=(12, 8))
sns.barplot(x='wage_bill', y='OCC_TITLE_MAJOR', data=df_wage_bill_major, palette="viridis")
plt.title('Total Wage Bill per Major Occupation Group')
plt.xlabel('Total Wage Bill (in billions)')
plt.ylabel('Major Occupation Group')
plt.grid(axis='x', linestyle='--', alpha=0.7)
def cell11():
# ───────────────────────────────────────────────────────────────
# 1. CUMULATIVE-DISTRIBUTION-FUNCTION (CDF) PREP
# ───────────────────────────────────────────────────────────────
def cdf(series):
s = series.sort_values().reset_index(drop=True)
return s.values, ((s.index + 1) / len(s)) * 100
x_lb , y_lb = cdf(atomic_tasks['lb_estimate_in_minutes'])
x_ub , y_ub = cdf(atomic_tasks['ub_estimate_in_minutes'])
x_mid, y_mid = cdf((atomic_tasks['ub_estimate_in_minutes'] + atomic_tasks['lb_estimate_in_minutes']) / 2)
# ───────────────────────────────────────────────────────────────
# 2. PLOTTING
# ───────────────────────────────────────────────────────────────
fig, ax = plt.subplots(figsize=(10, 6))
# horizontal reference lines every 10 %
for y_val in range(0, 101, 10):
ax.axhline(y_val, color=gray['100'], linewidth=.8, zorder=1)
# Plot Lower Bound CDF
ax.step(x_lb, y_lb,
where='post',
color=lime['300'], # Example: light blue for lower bound
linewidth=1.8,
linestyle='--',
zorder=2,
label='Lower bound estimate (CDF)')
# Plot Upper Bound CDF
ax.step(x_ub, y_ub,
where='post',
color=lime['900'], # Example: light orange/red for upper bound
linewidth=1.8,
linestyle=':',
zorder=3,
label='Upper bound estimate (CDF)')
# Plot Midpoint CDF (plotted last to be on top, or adjust zorder)
ax.step(x_mid, y_mid,
where='post',
color=lime['600'],
linewidth=2.2,
zorder=4, # Ensure it's on top of other lines if they overlap significantly
label='Mid-point estimate (CDF)')
# axes limits / scales
ax.set_ylim(0, 100)
ax.set_xscale('log')
# y-axis ➝ percent labels
ax.yaxis.set_major_formatter(mpl.ticker.PercentFormatter(decimals=0))
# move y-label to top-left (just inside plotting area)
ax.text(-0.06, 1.03,
"% of tasks with temporal coherence ≤ X",
ha='left', va='bottom',
transform=ax.transAxes,
fontsize=12, fontweight='semibold')
# custom x-ticks at human-friendly durations
ticks = [1, 5, 10, 30, 60, 120, 240, 480,
1440, 2880, 10080, 43200, 129600,
259200, 525600]
ticklabels = ['1 min', '5 min', '10 min', '30 min', '1 hour', '2 hours', '4 hours', '8 hours',
'1 day', '2 days', '1 week', '30 days',
'90 days', '180 days', '1 year']
# Vertical reference lines for x-ticks
for tick in ticks:
ax.axvline(tick, color=gray['300'], linewidth=.8, linestyle='--', zorder=1)
ax.set_xticks(ticks)
ax.set_xticklabels(ticklabels, rotation=45, ha='right')
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_edgecolor(gray['300'])
ax.spines['bottom'].set_edgecolor(gray['300'])
# legend
ax.legend(frameon=False, loc='lower right') # Keep 'lower right' or adjust as needed
ax.text(0.5, -0.3,
'Temporal coherence (X)',
ha='center', va='center',
transform=ax.transAxes,
fontsize=12, fontweight='semibold')

View file

@ -1,207 +0,0 @@
import logging
import re
import requests
import shutil
import sqlite3
import zipfile
from pathlib import Path
# Configure logging to provide feedback during the data setup process
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# --- Constants ---
# Using a data directory at the root of the project
DATA_DIR = Path("data")
# O*NET database details. We download the MySQL version and convert it to SQLite.
ONET_MYSQL_URL = "https://www.onetcenter.org/dl_files/database/db_29_3_mysql.zip"
DB_ZIP_PATH = DATA_DIR / "onet_mysql.zip"
DB_FILE_PATH = DATA_DIR / "onet.db"
EXTRACT_DIR = DATA_DIR / "onet_mysql_extracted"
# URLs for other required data files are in a separate text data archive.
ONET_TEXT_URL = "https://www.onetcenter.org/dl_files/database/db_29_3_text.zip"
TEXT_ZIP_PATH = DATA_DIR / "onet_text.zip"
TASK_RATINGS_PATH = DATA_DIR / "Task Ratings.txt"
DWA_REFERENCE_PATH = DATA_DIR / "DWA Reference.txt"
def setup_data_and_database():
"""
Main function to orchestrate the data setup.
It ensures the data directory exists, then downloads and sets up the O*NET database
and any other required data files.
"""
logging.info("Starting data and database setup...")
DATA_DIR.mkdir(exist_ok=True)
_setup_onet_database()
_download_additional_data()
logging.info("Data and database setup complete.")
def _setup_onet_database():
"""
Downloads the O*NET MySQL database, extracts it, and imports it into a
new SQLite database, following performance best practices from a shell script.
This method performs minimal text-based conversion of the MySQL dump to
make it compatible with SQLite before importing.
"""
if DB_FILE_PATH.exists():
logging.info("O*NET database already exists at %s. Skipping setup.", DB_FILE_PATH)
return
logging.info("O*NET database not found. Starting fresh setup.")
# Ensure the extraction directory is clean before use
if EXTRACT_DIR.exists():
shutil.rmtree(EXTRACT_DIR)
EXTRACT_DIR.mkdir()
try:
# 1. Download if necessary
if not DB_ZIP_PATH.exists():
logging.info("Downloading O*NET database from %s", ONET_MYSQL_URL)
_download_file(ONET_MYSQL_URL, DB_ZIP_PATH)
else:
logging.info("Using existing O*NET zip file at %s", DB_ZIP_PATH)
# 2. Extract
logging.info("Extracting O*NET database files to %s", EXTRACT_DIR)
with zipfile.ZipFile(DB_ZIP_PATH, 'r') as zip_ref:
zip_ref.extractall(EXTRACT_DIR)
# 3. Create new DB with performance PRAGMAs
logging.info("Creating new SQLite database with performance settings: %s", DB_FILE_PATH)
conn = sqlite3.connect(DB_FILE_PATH)
conn.executescript("""
PRAGMA journal_mode = OFF;
PRAGMA synchronous = 0;
PRAGMA cache_size = 1000000;
PRAGMA locking_mode = EXCLUSIVE;
PRAGMA temp_store = MEMORY;
""")
conn.close()
# 4. Combine all SQL files, convert, and import in a single transaction
logging.info("Combining and converting SQL files for single transaction import...")
sql_files = sorted(EXTRACT_DIR.rglob('*.sql'))
if not sql_files:
raise FileNotFoundError(f"No SQL files found in {EXTRACT_DIR}")
# Concatenate all files into one string
mysql_dump = "\n".join([sql_file.read_text(encoding='utf-8') for sql_file in sql_files])
# Minimal conversion for SQLite: remove backticks and ENGINE clauses
sqlite_dump = mysql_dump.replace('`', '')
sqlite_dump = re.sub(r'\) ENGINE=InnoDB.*?;', ');', sqlite_dump, flags=re.DOTALL)
full_script = f"BEGIN TRANSACTION;\n{sqlite_dump}\nCOMMIT;"
logging.info(f"Importing {len(sql_files)} SQL files into database...")
conn = sqlite3.connect(DB_FILE_PATH)
conn.executescript(full_script)
conn.close()
logging.info("Database populated successfully.")
# 5. Restore reliability settings and optimize
logging.info("Restoring reliability settings and optimizing database...")
conn = sqlite3.connect(DB_FILE_PATH)
conn.executescript("""
PRAGMA journal_mode = WAL;
PRAGMA synchronous = NORMAL;
PRAGMA locking_mode = NORMAL;
PRAGMA temp_store = DEFAULT;
PRAGMA foreign_keys = ON;
PRAGMA optimize;
""")
conn.execute("VACUUM;")
conn.close()
logging.info("Database setup and optimization complete.")
except Exception as e:
logging.error("Failed during database setup: %s", e, exc_info=True)
if DB_FILE_PATH.exists():
DB_FILE_PATH.unlink()
raise
finally:
# 6. Cleanup
logging.info("Cleaning up temporary files...")
if DB_ZIP_PATH.exists():
DB_ZIP_PATH.unlink()
if EXTRACT_DIR.exists():
shutil.rmtree(EXTRACT_DIR)
def _download_additional_data():
"""
Downloads and extracts supplementary data files from the O*NET text archive.
If the required text files already exist, this function does nothing.
"""
required_files = [TASK_RATINGS_PATH, DWA_REFERENCE_PATH]
if all(p.exists() for p in required_files):
logging.info("All required text data files already exist. Skipping download.")
return
logging.info("One or more text data files are missing. Downloading and extracting from archive...")
try:
_download_file(ONET_TEXT_URL, TEXT_ZIP_PATH)
logging.info("Unzipping text data archive...")
with zipfile.ZipFile(TEXT_ZIP_PATH, 'r') as zip_ref:
# Extract only the files we need, without creating subdirectories
for target_path in required_files:
if not target_path.exists():
# Find the corresponding file within the zip archive's directory structure
member_name = next((m for m in zip_ref.namelist() if m.endswith(target_path.name)), None)
if member_name:
with zip_ref.open(member_name) as source, open(target_path, 'wb') as target:
target.write(source.read())
logging.info("Extracted %s", target_path.name)
else:
logging.warning("Could not find %s in the text data archive.", target_path.name)
except requests.exceptions.RequestException as e:
logging.error("Failed to download O*NET text data archive: %s", e)
raise
except zipfile.BadZipFile as e:
logging.error("Failed to process the text data archive: %s", e)
raise
finally:
# Clean up the downloaded zip file
if TEXT_ZIP_PATH.exists():
TEXT_ZIP_PATH.unlink()
logging.info("Cleaned up downloaded text archive zip file.")
def _download_file(url, destination):
"""
Helper function to download a file from a URL, with streaming for large files.
"""
logging.info("Downloading from %s to %s", url, destination)
with requests.get(url, stream=True) as r:
r.raise_for_status()
with open(destination, 'wb') as f:
for chunk in r.iter_content(chunk_size=8192):
f.write(chunk)
logging.info("Download of %s complete.", destination.name)
def get_db_connection():
"""
Establishes and returns a connection to the SQLite database.
Returns None if the database file does not exist.
"""
if not DB_FILE_PATH.exists():
logging.error("Database file not found at %s. Run the setup process first.", DB_FILE_PATH)
return None
try:
conn = sqlite3.connect(DB_FILE_PATH)
return conn
except sqlite3.Error as e:
logging.error("Failed to connect to the database: %s", e)
return None
if __name__ == '__main__':
# This allows the data setup to be run directly from the command line,
# which is useful for initialization or debugging.
setup_data_and_database()

View file

@ -1,76 +0,0 @@
import importlib
import logging
import pkgutil
import shutil
from pathlib import Path
# The final destination for all generated outputs
DIST_DIR = Path("dist")
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
def create_all_outputs(processed_df):
"""
Dynamically discovers, imports, and runs all output generators.
This function iterates through all modules in the 'analysis.generators'
package. For each module, it assumes there is a 'generate(data)' function,
which it calls with the provided preprocessed DataFrame.
The generator function is expected to save its output to a temporary file
and return the path to that file. This function then moves the output
to the 'dist/' directory.
Args:
processed_df (pd.DataFrame): The fully preprocessed data to be used
by the generator functions.
"""
logging.info("Starting output generation...")
DIST_DIR.mkdir(exist_ok=True)
logging.info(f"Output directory is '{DIST_DIR.resolve()}'")
# Path to the generators package
from . import generators as generators_package
generators_path = generators_package.__path__
generators_prefix = generators_package.__name__ + "."
generated_files_count = 0
# Discover and run all modules in the generators package
for _, module_name, _ in pkgutil.iter_modules(generators_path, prefix=generators_prefix):
try:
logging.info(f"--- Running generator: {module_name} ---")
# Import the generator module
generator_module = importlib.import_module(module_name)
# Check if the module has the required 'generate' function
if not hasattr(generator_module, 'generate'):
logging.warning(f"Generator module {module_name} does not have a 'generate' function. Skipping.")
continue
# Call the generator function, passing in the preprocessed data
generator_func = getattr(generator_module, 'generate')
temp_output_path = generator_func(processed_df)
# If the generator returned a path, move the file to the dist directory
if temp_output_path and isinstance(temp_output_path, Path) and temp_output_path.exists():
# Sanitize the module name to create a valid filename
base_filename = module_name.split('.')[-1]
# Keep the original extension from the temp file
final_filename = base_filename + temp_output_path.suffix
final_output_path = DIST_DIR / final_filename
shutil.move(temp_output_path, final_output_path)
logging.info(f"Successfully generated '{final_output_path.name}'")
generated_files_count += 1
else:
logging.warning(f"Generator {module_name} did not return a valid output file path. Nothing was saved.")
except Exception as e:
logging.error(f"Failed to run generator {module_name}. Error: {e}", exc_info=True)
# Continue to the next generator
logging.info(f"--- Output generation complete. Total files generated: {generated_files_count} ---")

View file

@ -1,119 +0,0 @@
import seaborn as sns
import matplotlib.pyplot as plt
from pathlib import Path
import tempfile
import logging
import pandas as pd
import numpy as np
# Copied from other generators for modularity. This dictionary maps
# O*NET major occupation group codes to human-readable labels.
OCCUPATION_MAJOR_CODES = {
'11': 'Management',
'13': 'Business & Financial',
'15': 'Computer & Mathematical',
'17': 'Architecture & Engineering',
'19': 'Life, Physical, & Social Science',
'21': 'Community & Social Service',
'23': 'Legal',
'25': 'Education, Training, & Library',
'27': 'Arts, Design, & Media',
'29': 'Healthcare Practitioners',
'31': 'Healthcare Support',
'33': 'Protective Service',
'35': 'Food Preparation & Serving',
'37': 'Building & Grounds Maintenance',
'39': 'Personal Care & Service',
'41': 'Sales & Related',
'43': 'Office & Admin Support',
'45': 'Farming, Fishing, & Forestry',
'47': 'Construction & Extraction',
'49': 'Installation, Maintenance, & Repair',
'51': 'Production',
'53': 'Transportation & Material Moving',
'55': 'Military Specific',
}
def generate(processed_df: pd.DataFrame):
"""
Generates a scatter plot comparing lower vs. upper time estimates for tasks.
This corresponds to 'cell3' from the original analysis notebook. It helps
visualize the relationship and spread between the lower and upper bounds
of time estimates across different occupation groups.
Args:
processed_df (pd.DataFrame): The preprocessed data. Expected columns:
'lb_estimate_in_minutes',
'ub_estimate_in_minutes', 'onetsoc_major'.
Returns:
Path: The path to the generated temporary image file, or None on failure.
"""
logging.info("Generating plot of lower vs. upper time estimates...")
# --- Data Validation and Preparation ---
required_cols = ['lb_estimate_in_minutes', 'ub_estimate_in_minutes', 'onetsoc_major']
if not all(col in processed_df.columns for col in required_cols):
logging.error(f"Missing one or more required columns: {required_cols}. Cannot generate plot.")
return None
df = processed_df.copy()
# For log scaling, both lower and upper bounds must be positive.
df = df[(df['lb_estimate_in_minutes'] > 0) & (df['ub_estimate_in_minutes'] > 0)]
if df.empty:
logging.warning("No data with positive lower and upper estimates available to plot.")
return None
# Replace the major code with its readable label for the hue legend.
df['occupation_label'] = df['onetsoc_major'].map(OCCUPATION_MAJOR_CODES)
# --- Plotting ---
try:
plt.figure(figsize=(12, 10))
ax = sns.scatterplot(
data=df,
x='lb_estimate_in_minutes',
y='ub_estimate_in_minutes',
alpha=0.2,
edgecolor=None,
hue="occupation_label" # Use the labeled column for the legend
)
# Determine limits for the 45° reference line
# Use the maximum of both columns to create a square plot
max_val = df[['lb_estimate_in_minutes', 'ub_estimate_in_minutes']].max().max()
lims = (df[['lb_estimate_in_minutes', 'ub_estimate_in_minutes']].min().min(), max_val)
ax.plot(lims, lims, color='black', linestyle='--', linewidth=1, label='Upper = Lower')
# Add helper lines for constant ratios (2x, 10x, 100x)
for k in [2, 10, 100]:
ax.plot(lims, [k * l for l in lims],
linestyle=':', color='grey', linewidth=0.8, label=f'Upper = {k}x Lower')
ax.set(xscale='log', yscale='log', xlim=lims, ylim=lims)
ax.set_xlabel('Lower-bound Estimate (minutes, log scale)', fontsize=12)
ax.set_ylabel('Upper-bound Estimate (minutes, log scale)', fontsize=12)
ax.set_title('Lower vs. Upper Time Estimates for All Tasks', fontsize=16)
# Place the legend outside the plot to avoid obscuring data
ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left', title='Occupation / Ratio')
# --- File Saving ---
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "estimate_lower_vs_upper_bounds.png"
# Use bbox_inches='tight' to ensure the external legend is included in the saved image.
plt.savefig(temp_path, dpi=300, bbox_inches='tight')
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the plot: {e}", exc_info=True)
return None
finally:
plt.close()

View file

@ -1,86 +0,0 @@
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pathlib import Path
import tempfile
import logging
def generate(processed_df: pd.DataFrame):
"""
Generates a histogram of the log-ratio of upper to lower time estimates.
This corresponds to 'cell4' from the original analysis notebook. It shows
the distribution of how many times larger the upper estimate is compared
to the lower estimate.
Args:
processed_df (pd.DataFrame): The preprocessed data. Expected columns:
'lb_estimate_in_minutes',
'ub_estimate_in_minutes'.
Returns:
Path: The path to the generated temporary image file, or None on failure.
"""
logging.info("Generating distribution plot of estimate ratios...")
# --- Data Validation and Preparation ---
required_cols = ['lb_estimate_in_minutes', 'ub_estimate_in_minutes']
if not all(col in processed_df.columns for col in required_cols):
logging.error(f"Missing one or more required columns: {required_cols}. Cannot generate plot.")
return None
df = processed_df.copy()
# Calculate the ratio. We need to handle cases where the lower bound is zero.
# Replace lower bound of 0 with a small number to avoid division by zero, or filter them out.
# Here, we filter, as a ratio with a zero denominator is undefined.
df = df[df['lb_estimate_in_minutes'] > 0]
df['estimate_ratio'] = df['ub_estimate_in_minutes'] / df['lb_estimate_in_minutes']
# Replace infinite values (which can occur if ub is huge and lb is tiny) with NaN
# and drop rows with NaN or infinite ratios.
df.replace([np.inf, -np.inf], np.nan, inplace=True)
df.dropna(subset=['estimate_ratio'], inplace=True)
if df.empty:
logging.warning("No valid data available to plot the estimate ratio distribution.")
return None
# --- Plotting ---
try:
plt.figure(figsize=(10, 6))
# We plot the log10 of the ratio to better visualize the wide distribution
log_ratio = np.log10(df['estimate_ratio'])
sns.histplot(log_ratio, bins=60, kde=True)
# Add vertical lines for reference points
# log10(1) = 0, which is where upper bound equals lower bound
plt.axvline(x=0, color='black', linestyle='-', linewidth=1.5, label='1x (Upper = Lower)')
# A small ratio, e.g., 5% difference
plt.axvline(x=np.log10(1.05), color='orange', linestyle='--', linewidth=1, label='1.05x ratio')
# A 10x ratio
plt.axvline(x=np.log10(10), color='red', linestyle='--', linewidth=1, label='10x ratio')
plt.xlabel('log₁₀(Upper Estimate / Lower Estimate)', fontsize=12)
plt.ylabel('Number of Tasks', fontsize=12)
plt.title('Distribution of Time Estimate Ratios', fontsize=16)
plt.legend()
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.tight_layout()
# --- File Saving ---
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "estimate_ratio_distribution.png"
plt.savefig(temp_path, dpi=300)
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the plot: {e}", exc_info=True)
return None
finally:
plt.close()

View file

@ -1,135 +0,0 @@
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from pathlib import Path
import tempfile
import logging
# This mapping helps translate the O*NET 2-digit major group codes
# into human-readable labels for the plot's y-axis.
OCCUPATION_MAJOR_CODES = {
'11': 'Management',
'13': 'Business & Financial',
'15': 'Computer & Mathematical',
'17': 'Architecture & Engineering',
'19': 'Life, Physical, & Social Science',
'21': 'Community & Social Service',
'23': 'Legal',
'25': 'Education, Training, & Library',
'27': 'Arts, Design, & Media',
'29': 'Healthcare Practitioners',
'31': 'Healthcare Support',
'33': 'Protective Service',
'35': 'Food Preparation & Serving',
'37': 'Building & Grounds Maintenance',
'39': 'Personal Care & Service',
'41': 'Sales & Related',
'43': 'Office & Admin Support',
'45': 'Farming, Fishing, & Forestry',
'47': 'Construction & Extraction',
'49': 'Installation, Maintenance, & Repair',
'51': 'Production',
'53': 'Transportation & Material Moving',
'55': 'Military Specific',
}
def generate(processed_df: pd.DataFrame):
"""
Generates a heatmap of the median estimate ratio by occupation and task length quartile.
This corresponds to 'cell5' from the original analysis notebook. It shows
how the ratio between upper and lower time estimates varies across
different occupations and for tasks of different typical lengths (binned
into quartiles).
Args:
processed_df (pd.DataFrame): The preprocessed data. Expected columns:
'lb_estimate_in_minutes',
'ub_estimate_in_minutes', 'onetsoc_major'.
Returns:
Path: The path to the generated temporary image file, or None on failure.
"""
logging.info("Generating heatmap of estimate ratios by occupation and task length...")
# --- Data Validation and Preparation ---
required_cols = ['lb_estimate_in_minutes', 'ub_estimate_in_minutes', 'onetsoc_major']
if not all(col in processed_df.columns for col in required_cols):
logging.error(f"Missing one or more required columns: {required_cols}. Cannot generate plot.")
return None
df = processed_df.copy()
# Calculate the estimate ratio, handling division by zero and infinity
df = df[df['lb_estimate_in_minutes'] > 0]
df['estimate_ratio'] = df['ub_estimate_in_minutes'] / df['lb_estimate_in_minutes']
df.replace([np.inf, -np.inf], np.nan, inplace=True)
df.dropna(subset=['estimate_ratio'], inplace=True)
if df.empty:
logging.warning("No valid data available for the ratio heatmap.")
return None
# 1. Bin lower bounds into quartiles (Q1Q4)
# Using duplicates='drop' can help if there are many identical values
# which can make binning into quantiles fail.
try:
df['lb_q'] = pd.qcut(
df.lb_estimate_in_minutes,
q=4,
labels=['Q1 (Shortest)', 'Q2', 'Q3', 'Q4 (Longest)'],
duplicates='drop'
)
except ValueError as e:
logging.error(f"Could not bin data into quartiles: {e}. There might not be enough unique values.")
return None
# 2. Aggregate: median ratio per cell (occupation x task length quartile)
pivot = df.pivot_table(
index='onetsoc_major',
columns='lb_q',
values='estimate_ratio',
aggfunc='median'
)
# Map the index (onetsoc_major codes) to their corresponding readable labels
pivot.index = pivot.index.map(OCCUPATION_MAJOR_CODES)
pivot.dropna(inplace=True) # Drop occupations with no data in some quartiles for a cleaner plot
if pivot.empty:
logging.warning("Pivot table is empty after processing. Cannot generate heatmap.")
return None
# --- Plotting ---
try:
plt.figure(figsize=(12, 10))
sns.heatmap(
pivot,
cmap='RdYlGn_r', # Red-Yellow-Green (reversed), good for ratios centered around 1
center=2, # Center the colormap around a ratio of 2
annot=True, # Show the median values in the cells
fmt='.1f', # Format annotations to one decimal place
linewidths=.5,
cbar_kws={'label': 'Median Upper/Lower Estimate Ratio'}
)
plt.xlabel('Task Length (based on lower-bound quartile)', fontsize=12)
plt.ylabel('Occupation Major Group', fontsize=12)
plt.title('Typical Estimate Range Width by Occupation and Task Length', fontsize=16)
plt.tight_layout()
# --- File Saving ---
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "ratio_heatmap_by_occupation_and_task_length.png"
plt.savefig(temp_path, dpi=300)
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the heatmap: {e}", exc_info=True)
return None
finally:
plt.close()

View file

@ -1,161 +0,0 @@
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
import matplotlib.colors as mcolors
from pathlib import Path
import tempfile
import logging
# This mapping helps translate the O*NET 2-digit major group codes
# into human-readable labels for the plot's y-axis.
OCCUPATION_MAJOR_CODES = {
'11': 'Management',
'13': 'Business & Financial',
'15': 'Computer & Mathematical',
'17': 'Architecture & Engineering',
'19': 'Life, Physical, & Social Science',
'21': 'Community & Social Service',
'23': 'Legal',
'25': 'Education, Training, & Library',
'27': 'Arts, Design, & Media',
'29': 'Healthcare Practitioners',
'31': 'Healthcare Support',
'33': 'Protective Service',
'35': 'Food Preparation & Serving',
'37': 'Building & Grounds Maintenance',
'39': 'Personal Care & Service',
'41': 'Sales & Related',
'43': 'Office & Admin Support',
'45': 'Farming, Fishing, & Forestry',
'47': 'Construction & Extraction',
'49': 'Installation, Maintenance, & Repair',
'51': 'Production',
'53': 'Transportation & Material Moving',
'55': 'Military Specific',
}
# Define colors to match the original notebook's palette.
# These are standard hex codes for gray and lime shades.
BAR_COLORS = [
'#D1D5DB', # gray-300
'#84CC16', # lime-500
'#D9F99D', # lime-200
]
def _get_contrasting_text_color(bg_color_hex):
"""
Determines if black or white text provides better contrast against a given background color.
"""
try:
rgba = mcolors.to_rgba(bg_color_hex)
# Calculate luminance (Y) using the sRGB formula
luminance = 0.2126 * rgba[0] + 0.7152 * rgba[1] + 0.0722 * rgba[2]
return 'black' if luminance > 0.55 else 'white'
except ValueError:
return 'black' # Default to black if color is invalid
def generate(processed_df: pd.DataFrame):
"""
Generates a stacked bar chart breaking down tasks by remote status and estimability.
This corresponds to 'cell10' from the original analysis notebook. It shows,
for each occupation, the percentage of tasks that are not remote, remote and
estimable, or remote and not estimable.
Args:
processed_df (pd.DataFrame): The preprocessed data. Expected columns:
'onetsoc_major', 'remote_status', 'estimateable'.
Returns:
Path: The path to the generated temporary image file, or None on failure.
"""
logging.info("Generating task breakdown by occupation plot...")
# --- Data Validation ---
required_cols = ['onetsoc_major', 'remote_status', 'estimateable']
if not all(col in processed_df.columns for col in required_cols):
logging.error(f"Missing one or more required columns: {required_cols}. Cannot generate plot.")
return None
df = processed_df.copy()
# --- Data Summarization ---
summary_data = []
for code, label in OCCUPATION_MAJOR_CODES.items():
occ_df = df[df['onetsoc_major'] == code]
total_tasks = len(occ_df)
if total_tasks == 0:
continue
not_remote_count = len(occ_df[occ_df['remote_status'] != 'remote'])
remote_df = occ_df[occ_df['remote_status'] == 'remote']
remote_atomic_count = len(remote_df[remote_df['estimateable'] == 'ATOMIC'])
remote_ongoing_count = len(remote_df[remote_df['estimateable'] == 'ONGOING-CONSTRAINT'])
summary_data.append({
'occupation_label': label,
'count_not_remote': not_remote_count,
'count_remote_atomic': remote_atomic_count,
'count_remote_ongoing': remote_ongoing_count,
'total_tasks': total_tasks
})
if not summary_data:
logging.warning("No data available to generate the task breakdown plot.")
return None
summary_df = pd.DataFrame(summary_data)
# --- Percentage Calculation ---
summary_df['pct_not_remote'] = (summary_df['count_not_remote'] / summary_df['total_tasks']) * 100
summary_df['pct_remote_atomic'] = (summary_df['count_remote_atomic'] / summary_df['total_tasks']) * 100
summary_df['pct_remote_ongoing'] = (summary_df['count_remote_ongoing'] / summary_df['total_tasks']) * 100
plot_df = summary_df.set_index('occupation_label')[
['pct_not_remote', 'pct_remote_atomic', 'pct_remote_ongoing']
]
plot_df.columns = ['Not Remote', 'Remote & Estimable', 'Remote & Not Estimable']
plot_df = plot_df.sort_values(by='Not Remote', ascending=False)
# --- Plotting ---
try:
fig, ax = plt.subplots(figsize=(14, 10))
plot_df.plot(kind='barh', stacked=True, ax=ax, color=BAR_COLORS, width=0.8)
ax.set_xlabel("Percentage of Tasks", fontsize=12)
ax.set_ylabel("Occupation Major Group", fontsize=12)
ax.set_title("Task Breakdown by Occupation, Remote Status, and Estimability", fontsize=16, pad=20)
ax.xaxis.set_major_formatter(mtick.PercentFormatter())
ax.set_xlim(0, 100)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
# Add percentage labels inside each bar segment
for i, container in enumerate(ax.containers):
text_color = _get_contrasting_text_color(BAR_COLORS[i])
for patch in container.patches:
width = patch.get_width()
if width > 3: # Only label segments wider than 3%
x = patch.get_x() + width / 2
y = patch.get_y() + patch.get_height() / 2
ax.text(x, y, f"{width:.1f}%", ha='center', va='center',
fontsize=8, color=text_color, fontweight='medium')
ax.legend(title="Task Category", bbox_to_anchor=(1.02, 1), loc='upper left', frameon=False)
# --- File Saving ---
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "task_breakdown_by_occupation.png"
plt.savefig(temp_path, dpi=300, bbox_inches='tight')
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the plot: {e}", exc_info=True)
return None
finally:
plt.close()

View file

@ -1,74 +0,0 @@
import seaborn as sns
import matplotlib.pyplot as plt
from pathlib import Path
import tempfile
import logging
import pandas as pd
def generate(processed_df: pd.DataFrame):
"""
Generates a histogram of the task time estimate midpoints.
This generator corresponds to 'cell1' from the original analysis notebook.
It visualizes the distribution of the calculated midpoint of time estimates
for all tasks on a logarithmic scale to handle the wide range of values.
Args:
processed_df (pd.DataFrame): The preprocessed data, expected to contain
'lb_estimate_in_minutes' and
'ub_estimate_in_minutes' columns.
Returns:
Path: The path to the generated temporary image file, or None if
generation fails.
"""
logging.info("Generating task estimate distribution plot...")
# --- Data Validation and Preparation ---
required_cols = ['lb_estimate_in_minutes', 'ub_estimate_in_minutes']
if not all(col in processed_df.columns for col in required_cols):
logging.error(
f"Required columns {required_cols} not found in the DataFrame. "
"Cannot generate plot."
)
return None
# Create a copy to avoid modifying the original DataFrame
df = processed_df.copy()
# Calculate the midpoint from lower and upper bounds, as was done in the notebook
df['estimate_midpoint'] = (df['lb_estimate_in_minutes'] + df['ub_estimate_in_minutes']) / 2
# For log scaling, we must use positive values. Filter out any non-positive midpoints.
df = df[df['estimate_midpoint'] > 0]
if df.empty:
logging.warning("No data with positive estimate midpoints available to plot.")
return None
# --- Plotting ---
try:
plt.figure(figsize=(10, 6))
ax = sns.histplot(data=df, x='estimate_midpoint', log_scale=True)
ax.set_title('Distribution of Task Time Estimate Midpoints', fontsize=16)
ax.set_xlabel('Estimate Midpoint (minutes, log scale)', fontsize=12)
ax.set_ylabel('Number of Tasks', fontsize=12)
plt.tight_layout()
# --- File Saving ---
# Create a temporary file to save the plot. The orchestrator (`generate.py`)
# will move this to the final 'dist/' directory.
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "task_estimate_distribution.png"
plt.savefig(temp_path, dpi=300)
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the plot: {e}", exc_info=True)
return None
finally:
# Close the figure to free up memory, which is crucial when running many generators.
plt.close()

View file

@ -1,134 +0,0 @@
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from pathlib import Path
import tempfile
import logging
# Replicating the color palette from the original notebook for consistency.
# These appear to be inspired by Tailwind CSS colors.
GRAY_PALETTE = {
'100': '#F3F4F6',
'300': '#D1D5DB',
}
LIME_PALETTE = {
'300': '#D9F99D',
'600': '#A3E635', # A mid-tone lime
'900': '#4D7C0F', # A dark lime/green
}
def _calculate_cdf(series: pd.Series):
"""
Calculates the empirical Cumulative Distribution Function (CDF) for a series.
Returns the sorted values and their corresponding cumulative percentages.
"""
# Drop NA values and ensure the series is sorted
s = series.dropna().sort_values().reset_index(drop=True)
# Calculate cumulative percentage: (index + 1) / total_count
cdf_y = ((s.index + 1) / len(s)) * 100
return s.values, cdf_y
def generate(processed_df: pd.DataFrame):
"""
Generates a Cumulative Distribution Function (CDF) plot for task time estimates.
This corresponds to the second 'cell11' from the original notebook. It plots
the CDF for the lower-bound, upper-bound, and mid-point of time estimates,
showing the percentage of tasks that can be completed within a certain time.
Args:
processed_df (pd.DataFrame): The preprocessed data. Expected columns:
'lb_estimate_in_minutes',
'ub_estimate_in_minutes'.
Returns:
Path: The path to the generated temporary image file, or None on failure.
"""
logging.info("Generating temporal coherence CDF plot...")
# --- Data Validation and Preparation ---
required_cols = ['lb_estimate_in_minutes', 'ub_estimate_in_minutes']
if not all(col in processed_df.columns for col in required_cols):
logging.error(f"Missing one or more required columns: {required_cols}. Cannot generate plot.")
return None
df = processed_df.copy()
# Log scale requires positive values.
df = df[(df['lb_estimate_in_minutes'] > 0) & (df['ub_estimate_in_minutes'] > 0)]
if df.empty:
logging.warning("No data with positive estimates available to generate CDF plot.")
return None
# Calculate mid-point estimate
df['midpoint_estimate'] = (df['lb_estimate_in_minutes'] + df['ub_estimate_in_minutes']) / 2
# Prepare data for CDF plots
x_lb, y_lb = _calculate_cdf(df['lb_estimate_in_minutes'])
x_ub, y_ub = _calculate_cdf(df['ub_estimate_in_minutes'])
x_mid, y_mid = _calculate_cdf(df['midpoint_estimate'])
# --- Plotting ---
try:
fig, ax = plt.subplots(figsize=(12, 8))
# --- Grid and Reference Lines ---
# Horizontal reference lines for percentages
for y_val in range(0, 101, 10):
ax.axhline(y_val, color=GRAY_PALETTE['100'], linewidth=0.8, zorder=1)
# Vertical reference lines for human-friendly durations
ticks = [1, 5, 10, 30, 60, 120, 240, 480, 1440, 2880, 10080, 43200]
for tick in ticks:
ax.axvline(tick, color=GRAY_PALETTE['300'], linewidth=0.8, linestyle='--', zorder=1)
# --- CDF Plots ---
ax.step(x_lb, y_lb, where='post', color=LIME_PALETTE['300'], linewidth=1.8, linestyle='--', zorder=2, label='Lower-bound Estimate (CDF)')
ax.step(x_ub, y_ub, where='post', color=LIME_PALETTE['900'], linewidth=1.8, linestyle=':', zorder=3, label='Upper-bound Estimate (CDF)')
ax.step(x_mid, y_mid, where='post', color=LIME_PALETTE['600'], linewidth=2.2, zorder=4, label='Mid-point Estimate (CDF)')
# --- Axes Configuration ---
ax.set_ylim(0, 100)
ax.set_xscale('log')
# Custom x-ticks for durations
ticklabels = ['1 min', '5 min', '10 min', '30 min', '1 hr', '2 hrs', '4 hrs', '8 hrs', '1 day', '2 days', '1 week', '30 days']
ax.set_xticks(ticks)
ax.set_xticklabels(ticklabels, rotation=45, ha='right')
ax.minorticks_off() # Turn off minor ticks for clarity with custom grid
# Format y-axis as percentages
ax.yaxis.set_major_formatter(mpl.ticker.PercentFormatter(decimals=0))
# --- Spines and Labels ---
for spine in ['top', 'right']:
ax.spines[spine].set_visible(False)
for spine in ['left', 'bottom']:
ax.spines[spine].set_edgecolor(GRAY_PALETTE['300'])
# Use ax.text for more control over label placement than ax.set_ylabel/xlabel
ax.text(-0.07, 1.02, "% of tasks with duration ≤ X", transform=ax.transAxes,
fontsize=12, fontweight='semibold', va='bottom')
ax.text(0.5, -0.25, 'Task Duration (X)', transform=ax.transAxes,
fontsize=12, fontweight='semibold', ha='center')
ax.legend(frameon=False, loc='lower right')
fig.suptitle('Cumulative Distribution of Task Time Estimates', fontsize=16, y=0.96)
plt.tight_layout(rect=[0, 0, 1, 0.95]) # Adjust layout to make space for suptitle
# --- File Saving ---
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "temporal_coherence_cdf.png"
plt.savefig(temp_path, dpi=300, bbox_inches='tight')
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the CDF plot: {e}", exc_info=True)
return None
finally:
plt.close()

View file

@ -1,112 +0,0 @@
import seaborn as sns
import matplotlib.pyplot as plt
from pathlib import Path
import tempfile
import logging
import pandas as pd
# Based on O*NET SOC 2018 structure, this mapping helps translate
# the 2-digit major group codes into human-readable labels.
OCCUPATION_MAJOR_CODES = {
'11': 'Management',
'13': 'Business & Financial',
'15': 'Computer & Mathematical',
'17': 'Architecture & Engineering',
'19': 'Life, Physical, & Social Science',
'21': 'Community & Social Service',
'23': 'Legal',
'25': 'Education, Training, & Library',
'27': 'Arts, Design, & Media',
'29': 'Healthcare Practitioners',
'31': 'Healthcare Support',
'33': 'Protective Service',
'35': 'Food Preparation & Serving',
'37': 'Building & Grounds Maintenance',
'39': 'Personal Care & Service',
'41': 'Sales & Related',
'43': 'Office & Admin Support',
'45': 'Farming, Fishing, & Forestry',
'47': 'Construction & Extraction',
'49': 'Installation, Maintenance, & Repair',
'51': 'Production',
'53': 'Transportation & Material Moving',
'55': 'Military Specific',
}
def generate(processed_df: pd.DataFrame):
"""
Generates a box plot showing the spread of time-range estimates per occupation.
This corresponds to 'cell2' from the original analysis notebook. It visualizes
the distribution of the difference between upper and lower time estimates for
each major occupational group.
Args:
processed_df (pd.DataFrame): The preprocessed data. Expected columns:
'lb_estimate_in_minutes',
'ub_estimate_in_minutes', 'onetsoc_major'.
Returns:
Path: The path to the generated temporary image file, or None on failure.
"""
logging.info("Generating plot of time estimate spread by occupation...")
# --- Data Validation and Preparation ---
required_cols = ['lb_estimate_in_minutes', 'ub_estimate_in_minutes', 'onetsoc_major']
if not all(col in processed_df.columns for col in required_cols):
logging.error(f"Missing one or more required columns: {required_cols}. Cannot generate plot.")
return None
df = processed_df.copy()
# Calculate the estimate range.
df['estimate_range'] = df['ub_estimate_in_minutes'] - df['lb_estimate_in_minutes']
# For log scaling, we need positive values. Filter out any non-positive ranges.
df = df[df['estimate_range'] > 0]
if df.empty:
logging.warning("No data with a positive estimate range available to plot.")
return None
# Sort by the major code to ensure a consistent plot order
df = df.sort_values('onetsoc_major')
# --- Plotting ---
try:
plt.figure(figsize=(14, 10))
ax = sns.boxplot(
data=df,
x='onetsoc_major',
y='estimate_range',
showfliers=False # Outliers are excluded for a clearer view of the main distribution
)
plt.yscale('log') # The long tail of the data makes a log scale more readable
plt.xlabel('Occupation Major Group', fontsize=12)
plt.ylabel('Time Estimate Range (upper - lower, in minutes, log scale)', fontsize=12)
plt.title('Spread of Time-Range Estimates by Occupation', fontsize=16)
# Replace numeric x-tick labels (e.g., '11', '15') with meaningful text labels
ax.set_xticklabels(
[OCCUPATION_MAJOR_CODES.get(code.get_text(), code.get_text()) for code in ax.get_xticklabels()],
rotation=60,
ha='right' # Align rotated labels correctly
)
plt.tight_layout()
# --- File Saving ---
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "time_estimate_spread_by_occupation.png"
plt.savefig(temp_path, dpi=300, bbox_inches='tight')
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the plot: {e}", exc_info=True)
return None
finally:
plt.close()

View file

@ -1,150 +0,0 @@
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import pandas as pd
from pathlib import Path
import tempfile
import logging
# Assuming data.py is in the same package and provides this function
from ..data import get_db_connection
# This mapping helps translate the O*NET 2-digit major group codes
# into human-readable labels for the plot's y-axis.
OCCUPATION_MAJOR_CODES = {
'11': 'Management',
'13': 'Business & Financial',
'15': 'Computer & Mathematical',
'17': 'Architecture & Engineering',
'19': 'Life, Physical, & Social Science',
'21': 'Community & Social Service',
'23': 'Legal',
'25': 'Education, Training, & Library',
'27': 'Arts, Design, & Media',
'29': 'Healthcare Practitioners',
'31': 'Healthcare Support',
'33': 'Protective Service',
'35': 'Food Preparation & Serving',
'37': 'Building & Grounds Maintenance',
'39': 'Personal Care & Service',
'41': 'Sales & Related',
'43': 'Office & Admin Support',
'45': 'Farming, Fishing, & Forestry',
'47': 'Construction & Extraction',
'49': 'Installation, Maintenance, & Repair',
'51': 'Production',
'53': 'Transportation & Material Moving',
'55': 'Military Specific',
}
def generate(processed_df: pd.DataFrame):
"""
Generates a bar plot of the total wage bill per major occupation group.
This corresponds to the first 'cell11' from the original analysis notebook.
It calculates the total wage bill (Total Employment * Annual Mean Wage) for
each occupation and aggregates it by major occupation group. This generator
loads its data directly from the O*NET database.
Args:
processed_df (pd.DataFrame): The preprocessed data (not used in this generator,
but required by the function signature).
Returns:
Path: The path to the generated temporary image file, or None on failure.
"""
logging.info("Generating plot of total wage bill by occupation...")
conn = None
try:
# --- Data Loading ---
# This generator needs specific data that is not in the main preprocessed_df.
# It loads occupational employment and wage data directly from the database.
conn = get_db_connection()
if conn is None:
raise ConnectionError("Could not get database connection.")
# This data is stored in a long format in the `occupation_level_metadata` table.
# We need to query this table and pivot it to get employment and wage columns.
query = "SELECT onetsoc_code, item, response FROM occupation_level_metadata WHERE item IN ('Employment', 'Annual Mean Wage')"
try:
df_meta = pd.read_sql_query(query, conn)
# Pivot the table to create 'Employment' and 'Annual Mean Wage' columns
df_oesm = df_meta.pivot(index='onetsoc_code', columns='item', values='response').reset_index()
logging.info("Pivoted occupation metadata. Columns are: %s", df_oesm.columns.tolist())
# Rename for consistency with the original notebook's code
df_oesm.rename(columns={
'onetsoc_code': 'OCC_CODE',
'Employment': 'TOT_EMP',
'Annual Mean Wage': 'A_MEAN'
}, inplace=True)
except (pd.io.sql.DatabaseError, KeyError) as e:
logging.error(f"Failed to query or pivot occupation metadata: {e}", exc_info=True)
return None
# --- Data Preparation ---
# Create a 'major group' code from the first two digits of the SOC code
df_oesm['onetsoc_major'] = df_oesm['OCC_CODE'].str[:2]
# Ensure wage and employment columns are numeric, coercing errors to NaN
df_oesm['TOT_EMP'] = pd.to_numeric(df_oesm['TOT_EMP'], errors='coerce')
df_oesm['A_MEAN'] = pd.to_numeric(df_oesm['A_MEAN'], errors='coerce')
# Drop rows with missing data in critical columns
df_oesm.dropna(subset=['TOT_EMP', 'A_MEAN', 'onetsoc_major'], inplace=True)
# Calculate the wage bill for each occupation
df_oesm['wage_bill'] = df_oesm['TOT_EMP'] * df_oesm['A_MEAN']
# Aggregate the wage bill by major occupation group
df_wage_bill_major = df_oesm.groupby('onetsoc_major')['wage_bill'].sum().reset_index()
# Map the major codes to readable titles for plotting
df_wage_bill_major['OCC_TITLE_MAJOR'] = df_wage_bill_major['onetsoc_major'].map(OCCUPATION_MAJOR_CODES)
df_wage_bill_major.dropna(subset=['OCC_TITLE_MAJOR'], inplace=True) # Drop military/unmapped codes
# Sort by wage bill for a more informative plot
df_wage_bill_major = df_wage_bill_major.sort_values('wage_bill', ascending=False)
if df_wage_bill_major.empty:
logging.warning("No data available to generate the wage bill plot.")
return None
# --- Plotting ---
plt.figure(figsize=(12, 10))
ax = sns.barplot(x='wage_bill', y='OCC_TITLE_MAJOR', data=df_wage_bill_major, palette="viridis", orient='h')
ax.set_title('Total Wage Bill per Major Occupation Group', fontsize=16, pad=15)
ax.set_xlabel('Total Wage Bill (in USD)', fontsize=12)
ax.set_ylabel('Major Occupation Group', fontsize=12)
ax.grid(axis='x', linestyle='--', alpha=0.7)
# Format the x-axis to be more readable (e.g., "$2.0T" for trillions)
def format_billions(x, pos):
if x >= 1e12:
return f'${x*1e-12:.1f}T'
if x >= 1e9:
return f'${x*1e-9:.0f}B'
return f'${x*1e-6:.0f}M'
ax.xaxis.set_major_formatter(mticker.FuncFormatter(format_billions))
plt.tight_layout()
# --- File Saving ---
temp_dir = tempfile.gettempdir()
temp_path = Path(temp_dir) / "wage_bill_by_occupation.png"
plt.savefig(temp_path, dpi=300)
logging.info(f"Successfully saved plot to temporary file: {temp_path}")
return temp_path
except Exception as e:
logging.error(f"An error occurred while generating the wage bill plot: {e}", exc_info=True)
return None
finally:
plt.close()
if conn:
conn.close()

View file

@ -1,64 +0,0 @@
import logging
import sys
# Since this file is inside the 'analysis' package, we use relative imports
# to access the other modules within the same package.
from . import data
from . import preprocess
from . import generate
# Configure logging for the entire application.
# This setup will apply to loggers in data, preprocess, and generate modules as well.
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
stream=sys.stdout
)
def main():
"""
The main entry point for the entire analysis pipeline.
This function orchestrates the three main stages of the analysis:
1. Data Setup: Downloads and prepares the necessary raw data and database.
2. Preprocessing: Cleans, enriches, and transforms the raw data into an
analysis-ready DataFrame.
3. Output Generation: Runs all registered generators to produce figures,
tables, and other outputs, saving them to the 'dist/' directory.
"""
logger = logging.getLogger(__name__)
logger.info("=================================================")
logger.info(" STARTING ECONTAI ANALYSIS PIPELINE ")
logger.info("=================================================")
try:
# Stage 1: Set up the data and database
logger.info("--- STAGE 1: DATA SETUP ---")
data.setup_data_and_database()
logger.info("--- DATA SETUP COMPLETE ---")
# Stage 2: Run the preprocessing pipeline
logger.info("--- STAGE 2: PREPROCESSING ---")
processed_dataframe = preprocess.run_preprocessing()
logger.info("--- PREPROCESSING COMPLETE ---")
# Stage 3: Generate all outputs
logger.info("--- STAGE 3: OUTPUT GENERATION ---")
generate.create_all_outputs(processed_dataframe)
logger.info("--- OUTPUT GENERATION COMPLETE ---")
logger.info("=================================================")
logger.info(" ANALYSIS PIPELINE COMPLETED SUCCESSFULLY ")
logger.info("=================================================")
except Exception as e:
logger.critical("An unrecoverable error occurred during the pipeline execution.", exc_info=True)
# Exit with a non-zero status code to indicate failure, which is useful for automation.
sys.exit(1)
# This allows the script to be run from the command line using `python -m analysis.main`.
# The `-m` flag is important because it adds the parent directory to the Python path,
# allowing the relative imports (e.g., `from . import data`) to work correctly.
if __name__ == '__main__':
main()

View file

@ -1,160 +0,0 @@
import logging
import pandas as pd
import numpy as np
from scipy.stats import median_abs_deviation
from .data import get_db_connection
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
def _convert_to_minutes(level: float) -> float:
"""
Converts O*NET 'Frequency' scale values (levels) to estimated minutes per day.
This logic is derived from the `preprocessing_time_estimates` function
in the original analysis notebook.
"""
if pd.isna(level):
return 0
# This mapping is an interpretation of the O*NET frequency scale.
return {
1: 0, # Yearly or less
2: 2, # Several times a year
3: 10, # Several times a month
4: 30, # Several times a week
5: 120, # Daily
6: 240, # Several times a day
7: 480, # Hourly or more
}.get(int(level), 0)
def _mad_z_score(series: pd.Series) -> pd.Series:
"""
Calculates the robust Z-score using Median Absolute Deviation (MAD).
This function is derived from 'cell7' of the original analysis.
"""
if series.isnull().all():
return pd.Series([np.nan] * len(series), index=series.index)
median = series.median()
# scale='normal' makes MAD comparable to the standard deviation for a normal distribution.
mad = median_abs_deviation(series.dropna(), scale='normal')
if mad == 0:
return pd.Series([np.nan] * len(series), index=series.index)
return (series - median) / mad
def run_preprocessing() -> pd.DataFrame:
"""
Main orchestrator for the preprocessing pipeline.
This function faithfully reproduces the data transformation pipeline from the
original `analysis.py` script, including the `preprocessing_time_estimates`
and cell-specific data manipulations.
Returns:
pd.DataFrame: A fully preprocessed DataFrame ready for the generators.
"""
logging.info("Starting data preprocessing...")
conn = None
try:
conn = get_db_connection()
if conn is None:
raise ConnectionError("Could not establish database connection.")
# --- 1. Load Data from Database ---
# Fetch all necessary tables to build the initial DataFrame.
logging.info("Loading data from O*NET database...")
task_ratings_df = pd.read_sql_query("SELECT * FROM task_ratings", conn)
task_statements_df = pd.read_sql_query("SELECT * FROM task_statements", conn)
occupations_df = pd.read_sql_query("SELECT * FROM occupation_data", conn)
# --- 2. Initial Merge ---
# Merge the tables to create a comprehensive base DataFrame.
# Merging on both 'onetsoc_code' and 'task_id' is crucial to avoid
# creating duplicate columns from the overlapping 'onetsoc_code'.
logging.info("Merging base tables...")
tasks_df = pd.merge(task_ratings_df, task_statements_df, on=['onetsoc_code', 'task_id'])
tasks_df = pd.merge(tasks_df, occupations_df, on='onetsoc_code')
# --- 3. Create "Atomic Tasks" and Time Estimates (from `preprocessing_time_estimates`) ---
# This is the core of the analysis, focusing on tasks with frequency ratings.
logging.info("Filtering for 'atomic tasks' (scale_id='FR') and calculating time estimates...")
# Strip whitespace from scale_id to ensure the filter works correctly.
tasks_df['scale_id'] = tasks_df['scale_id'].str.strip()
atomic_tasks = tasks_df[tasks_df['scale_id'] == 'FR'].copy()
# Convert frequency confidence intervals into minutes/day
atomic_tasks['lb_estimate_in_minutes'] = atomic_tasks['lower_ci_bound'].apply(_convert_to_minutes)
atomic_tasks['ub_estimate_in_minutes'] = atomic_tasks['upper_ci_bound'].apply(_convert_to_minutes)
atomic_tasks['estimate_midpoint'] = (atomic_tasks['lb_estimate_in_minutes'] + atomic_tasks['ub_estimate_in_minutes']) / 2
# --- 4. Add Derived Columns for Analysis (from `cell` logic) ---
logging.info("Adding derived columns for analysis...")
# Add `onetsoc_major` for grouping by occupation category
atomic_tasks['onetsoc_major'] = atomic_tasks['onetsoc_code'].str[:2]
# Calculate estimate_range and estimate_ratio used in several plots
atomic_tasks['estimate_range'] = atomic_tasks['ub_estimate_in_minutes'] - atomic_tasks['lb_estimate_in_minutes']
# To calculate ratio, ensure lower bound is positive to avoid division by zero
lb_positive = atomic_tasks['lb_estimate_in_minutes'] > 0
atomic_tasks['estimate_ratio'] = np.nan
atomic_tasks.loc[lb_positive, 'estimate_ratio'] = atomic_tasks['ub_estimate_in_minutes'] / atomic_tasks['lb_estimate_in_minutes']
# --- 5. Calculate Outlier Scores (from `cell6` and `cell7`) ---
logging.info("Calculating standard and robust Z-scores for outlier detection...")
# Standard Z-score
grouped_stats = atomic_tasks.groupby('onetsoc_code')['estimate_midpoint'].agg(['mean', 'std'])
atomic_tasks = atomic_tasks.merge(grouped_stats, on='onetsoc_code', how='left')
# Calculate Z-score, avoiding division by zero if std is 0
non_zero_std = atomic_tasks['std'].notna() & (atomic_tasks['std'] != 0)
atomic_tasks['z_score'] = np.nan
atomic_tasks.loc[non_zero_std, 'z_score'] = \
(atomic_tasks.loc[non_zero_std, 'estimate_midpoint'] - atomic_tasks.loc[non_zero_std, 'mean']) / atomic_tasks.loc[non_zero_std, 'std']
# Robust Z-score (using MAD)
atomic_tasks['robust_z_score'] = atomic_tasks.groupby('onetsoc_code')['estimate_midpoint'].transform(_mad_z_score)
# --- 6. Prepare for other generators ---
# NOTE: The data for the 'task_breakdown_by_occupation' generator, specifically
# the 'remote_status' and 'estimateable' columns, is not available in the O*NET
# database. This data was likely loaded from a separate file (e.g., 'tasks_clean.parquet')
# in the original notebook. For now, we will add placeholder columns.
atomic_tasks['remote_status'] = 'unknown'
atomic_tasks['estimateable'] = 'unknown'
logging.info("Data preprocessing complete.")
return atomic_tasks
except Exception as e:
logging.error("An error occurred during preprocessing: %s", e, exc_info=True)
# Return an empty DataFrame on failure to prevent downstream errors
return pd.DataFrame()
finally:
if conn:
conn.close()
logging.info("Database connection closed.")
if __name__ == '__main__':
# This allows the preprocessing to be run directly for testing or debugging.
# Note: Requires data to be set up first by running data.py.
try:
processed_data = run_preprocessing()
if not processed_data.empty:
print("Preprocessing successful. DataFrame shape:", processed_data.shape)
print("Columns:", processed_data.columns.tolist())
print(processed_data.head())
# Save to a temporary file to inspect the output
output_path = "temp_preprocessed_data.csv"
processed_data.to_csv(output_path, index=False)
print(f"Sample output saved to {output_path}")
else:
print("Preprocessing failed or resulted in an empty DataFrame.")
except (FileNotFoundError, ConnectionError) as e:
logging.error("Failed to run preprocessing: %s", e)

View file

@ -1,22 +0,0 @@
# Presentation
## Notebooks
- [data enrichment](data_enrichment.ipynb) - contains the code to gather things from the O\*NET data, BLS's OEWS database (unused for now), Barnett's data...
- [prompt evaluation](evaluate_llm_time_estimations.ipynb) - the playground used to evaluate change in hyperparameters (system prompt, user prompt, schema, model...)
- [analysis](analysis.ipynb) - code to generate the graphs in the paper
- [legacy](legacy.ipynb) - if there are some missing pieces, it's worth looking in there.
## Running the non-notebook code
To re-run everything, you need python and uv up and running, if you use have nix installed, run
```bash
nix develop .#impure
```
and then `uv run ...` as requested in the notebooks.
If some things are missing, email <dorn@xfe.li>, I'm usually reactive.
Copy `.env.example` to `.env` and fill in OPENAI_API_KEY. The total run and experiments cost less than <10$.

File diff suppressed because one or more lines are too long

File diff suppressed because it is too large Load diff

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -1,628 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 86,
"id": "beace815-b5ae-44a4-a81c-a7f82cb66296",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2K\u001b[2mResolved \u001b[1m118 packages\u001b[0m \u001b[2min 386ms\u001b[0m\u001b[0m \u001b[0m\n",
"\u001b[2K\u001b[2mPrepared \u001b[1m2 packages\u001b[0m \u001b[2min 124ms\u001b[0m\u001b[0m \n",
"\u001b[2K\u001b[2mInstalled \u001b[1m2 packages\u001b[0m \u001b[2min 5ms\u001b[0m\u001b[0m \u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1met-xmlfile\u001b[0m\u001b[2m==2.0.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mopenpyxl\u001b[0m\u001b[2m==3.1.5\u001b[0m\n"
]
}
],
"source": [
"!uv add pandas requests openai dotenv openpyxl"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "941d511f-ad72-4306-bbab-52127583e513",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import dotenv\n",
"import openai\n",
"import sqlite3\n",
"import pandas as pd\n",
"\n",
"dotenv.load_dotenv() # Copy .env.example to .env and fill in the blanks\n",
"\n",
"oai_token = os.getenv(\"OPENAI_API_KEY\")\n",
"\n",
"oai = openai.OpenAI(api_key=oai_token)\n",
"onet = sqlite3.connect(\"onet.database\") # Run ./create_onet_database.sh to create it\n",
"# This dataset comes from https://epoch.ai/gradient-updates/consequences-of-automating-remote-work\n",
"# It contains labels for whethere a O*NET task can be done remotely or not (labeled by GPT-4o)\n",
"# You can download it here: https://drive.google.com/file/d/1GrHhuYIgaCCgo99dZ_40BWraz-fzo76r/view?usp=sharing\n",
"df_remote_status = pd.read_csv(\"epoch_task_data.csv\")\n",
"\n",
"# BLS OEWS: https://www.bls.gov/oes/special-requests/oesm23nat.zip\n",
"df_oesm = pd.read_excel(\"oesm23national.xlsx\")\n",
"\n",
"# Run uv run enrich_task_ratings.py to get this file (trs = Task RatingS)\n",
"df_enriched_trs = pd.read_json(\"task_ratings_enriched.json\")"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "a5351f8b-c2ad-4d3e-af4a-992f539a6064",
"metadata": {},
"outputs": [],
"source": [
"FREQUENCY_MAP = {\n",
" 'frequency_category_1': \"Yearly or less\",\n",
" 'frequency_category_2': \"More than yearly\",\n",
" 'frequency_category_3': \"More than monthly\",\n",
" 'frequency_category_4': \"More than weekly\",\n",
" 'frequency_category_5': \"Daily\",\n",
" 'frequency_category_6': \"Several times daily\",\n",
" 'frequency_category_7': \"Hourly or more\"\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "8b2ab22a-afab-41f9-81a3-48eab261b568",
"metadata": {},
"outputs": [],
"source": [
"background_prompt = '''\n",
"Estimate the typical duration to complete *one instance* of the following job task from the moment a person starts to work on it to the last moment the person will need to keep it in mind\n",
"\n",
"Take into account that there might be delays between the steps to complete the task, which would lengthen the estimate.\n",
"\n",
"Output a range with the format [duration A] - [duration B] where [duration A] and [duration B] correspond to one of the durations below:\n",
"- less than 30 minutes\n",
"- 30 minutes\n",
"- 1 hour\n",
"- 4 hours\n",
"- 8 hours\n",
"- 16 hours\n",
"- 3 days\n",
"- 1 week\n",
"- 3 weeks\n",
"- 6 weeks\n",
"- 3 months\n",
"- 6 months\n",
"- 1 year\n",
"- 3 years\n",
"- more than 3 year\n",
"\n",
"**Do not output anything besides the range**\n",
"'''"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "d2e4a855-f327-4b3d-ad0b-ed997e720639",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>onetsoc_code</th>\n",
" <th>task_id</th>\n",
" <th>task</th>\n",
" <th>occupation_title</th>\n",
" <th>occupation_description</th>\n",
" <th>frequency_category_1</th>\n",
" <th>frequency_category_2</th>\n",
" <th>frequency_category_3</th>\n",
" <th>frequency_category_4</th>\n",
" <th>frequency_category_5</th>\n",
" <th>frequency_category_6</th>\n",
" <th>frequency_category_7</th>\n",
" <th>importance_average</th>\n",
" <th>relevance_average</th>\n",
" <th>OCC_CODE</th>\n",
" <th>TOT_EMP</th>\n",
" <th>H_MEAN</th>\n",
" <th>A_MEAN</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>11-1011.00</td>\n",
" <td>8823</td>\n",
" <td>Direct or coordinate an organization's financi...</td>\n",
" <td>Chief Executives</td>\n",
" <td>Determine and formulate policies and provide o...</td>\n",
" <td>5.92</td>\n",
" <td>15.98</td>\n",
" <td>29.68</td>\n",
" <td>21.18</td>\n",
" <td>19.71</td>\n",
" <td>4.91</td>\n",
" <td>2.63</td>\n",
" <td>4.52</td>\n",
" <td>74.44</td>\n",
" <td>11-1011</td>\n",
" <td>211230.0</td>\n",
" <td>124.47</td>\n",
" <td>258900</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>11-1011.00</td>\n",
" <td>8824</td>\n",
" <td>Confer with board members, organization offici...</td>\n",
" <td>Chief Executives</td>\n",
" <td>Determine and formulate policies and provide o...</td>\n",
" <td>1.42</td>\n",
" <td>14.44</td>\n",
" <td>27.31</td>\n",
" <td>25.52</td>\n",
" <td>26.88</td>\n",
" <td>2.52</td>\n",
" <td>1.90</td>\n",
" <td>4.32</td>\n",
" <td>81.71</td>\n",
" <td>11-1011</td>\n",
" <td>211230.0</td>\n",
" <td>124.47</td>\n",
" <td>258900</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>11-1011.00</td>\n",
" <td>8827</td>\n",
" <td>Prepare budgets for approval, including those ...</td>\n",
" <td>Chief Executives</td>\n",
" <td>Determine and formulate policies and provide o...</td>\n",
" <td>15.50</td>\n",
" <td>38.21</td>\n",
" <td>32.73</td>\n",
" <td>5.15</td>\n",
" <td>5.25</td>\n",
" <td>0.19</td>\n",
" <td>2.98</td>\n",
" <td>4.30</td>\n",
" <td>93.41</td>\n",
" <td>11-1011</td>\n",
" <td>211230.0</td>\n",
" <td>124.47</td>\n",
" <td>258900</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>11-1011.00</td>\n",
" <td>8826</td>\n",
" <td>Direct, plan, or implement policies, objective...</td>\n",
" <td>Chief Executives</td>\n",
" <td>Determine and formulate policies and provide o...</td>\n",
" <td>3.03</td>\n",
" <td>17.33</td>\n",
" <td>20.30</td>\n",
" <td>18.10</td>\n",
" <td>33.16</td>\n",
" <td>2.01</td>\n",
" <td>6.07</td>\n",
" <td>4.24</td>\n",
" <td>97.79</td>\n",
" <td>11-1011</td>\n",
" <td>211230.0</td>\n",
" <td>124.47</td>\n",
" <td>258900</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>11-1011.00</td>\n",
" <td>8834</td>\n",
" <td>Prepare or present reports concerning activiti...</td>\n",
" <td>Chief Executives</td>\n",
" <td>Determine and formulate policies and provide o...</td>\n",
" <td>1.98</td>\n",
" <td>14.06</td>\n",
" <td>42.60</td>\n",
" <td>21.24</td>\n",
" <td>13.18</td>\n",
" <td>6.24</td>\n",
" <td>0.70</td>\n",
" <td>4.17</td>\n",
" <td>92.92</td>\n",
" <td>11-1011</td>\n",
" <td>211230.0</td>\n",
" <td>124.47</td>\n",
" <td>258900</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>17634</th>\n",
" <td>53-7121.00</td>\n",
" <td>12807</td>\n",
" <td>Unload cars containing liquids by connecting h...</td>\n",
" <td>Tank Car, Truck, and Ship Loaders</td>\n",
" <td>Load and unload chemicals and bulk solids, suc...</td>\n",
" <td>6.05</td>\n",
" <td>29.21</td>\n",
" <td>6.88</td>\n",
" <td>13.95</td>\n",
" <td>27.65</td>\n",
" <td>7.93</td>\n",
" <td>8.34</td>\n",
" <td>4.08</td>\n",
" <td>64.04</td>\n",
" <td>53-7121</td>\n",
" <td>11400.0</td>\n",
" <td>29.1</td>\n",
" <td>60530</td>\n",
" </tr>\n",
" <tr>\n",
" <th>17635</th>\n",
" <td>53-7121.00</td>\n",
" <td>12804</td>\n",
" <td>Clean interiors of tank cars or tank trucks, u...</td>\n",
" <td>Tank Car, Truck, and Ship Loaders</td>\n",
" <td>Load and unload chemicals and bulk solids, suc...</td>\n",
" <td>1.47</td>\n",
" <td>6.33</td>\n",
" <td>21.70</td>\n",
" <td>25.69</td>\n",
" <td>32.35</td>\n",
" <td>12.47</td>\n",
" <td>0.00</td>\n",
" <td>4.02</td>\n",
" <td>44.33</td>\n",
" <td>53-7121</td>\n",
" <td>11400.0</td>\n",
" <td>29.1</td>\n",
" <td>60530</td>\n",
" </tr>\n",
" <tr>\n",
" <th>17636</th>\n",
" <td>53-7121.00</td>\n",
" <td>12803</td>\n",
" <td>Lower gauge rods into tanks or read meters to ...</td>\n",
" <td>Tank Car, Truck, and Ship Loaders</td>\n",
" <td>Load and unload chemicals and bulk solids, suc...</td>\n",
" <td>4.52</td>\n",
" <td>1.76</td>\n",
" <td>4.65</td>\n",
" <td>17.81</td>\n",
" <td>37.42</td>\n",
" <td>23.31</td>\n",
" <td>10.55</td>\n",
" <td>3.88</td>\n",
" <td>65.00</td>\n",
" <td>53-7121</td>\n",
" <td>11400.0</td>\n",
" <td>29.1</td>\n",
" <td>60530</td>\n",
" </tr>\n",
" <tr>\n",
" <th>17637</th>\n",
" <td>53-7121.00</td>\n",
" <td>12805</td>\n",
" <td>Operate conveyors and equipment to transfer gr...</td>\n",
" <td>Tank Car, Truck, and Ship Loaders</td>\n",
" <td>Load and unload chemicals and bulk solids, suc...</td>\n",
" <td>6.97</td>\n",
" <td>12.00</td>\n",
" <td>2.52</td>\n",
" <td>5.90</td>\n",
" <td>35.48</td>\n",
" <td>22.08</td>\n",
" <td>15.05</td>\n",
" <td>3.87</td>\n",
" <td>47.90</td>\n",
" <td>53-7121</td>\n",
" <td>11400.0</td>\n",
" <td>29.1</td>\n",
" <td>60530</td>\n",
" </tr>\n",
" <tr>\n",
" <th>17638</th>\n",
" <td>53-7121.00</td>\n",
" <td>12810</td>\n",
" <td>Perform general warehouse activities, such as ...</td>\n",
" <td>Tank Car, Truck, and Ship Loaders</td>\n",
" <td>Load and unload chemicals and bulk solids, suc...</td>\n",
" <td>5.91</td>\n",
" <td>10.85</td>\n",
" <td>6.46</td>\n",
" <td>14.46</td>\n",
" <td>34.14</td>\n",
" <td>16.39</td>\n",
" <td>11.78</td>\n",
" <td>3.53</td>\n",
" <td>47.84</td>\n",
" <td>53-7121</td>\n",
" <td>11400.0</td>\n",
" <td>29.1</td>\n",
" <td>60530</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"<p>17639 rows × 18 columns</p>\n",
"</div>"
],
"text/plain": [
" onetsoc_code task_id \\\n",
"0 11-1011.00 8823 \n",
"1 11-1011.00 8824 \n",
"2 11-1011.00 8827 \n",
"3 11-1011.00 8826 \n",
"4 11-1011.00 8834 \n",
"... ... ... \n",
"17634 53-7121.00 12807 \n",
"17635 53-7121.00 12804 \n",
"17636 53-7121.00 12803 \n",
"17637 53-7121.00 12805 \n",
"17638 53-7121.00 12810 \n",
"\n",
" task \\\n",
"0 Direct or coordinate an organization's financi... \n",
"1 Confer with board members, organization offici... \n",
"2 Prepare budgets for approval, including those ... \n",
"3 Direct, plan, or implement policies, objective... \n",
"4 Prepare or present reports concerning activiti... \n",
"... ... \n",
"17634 Unload cars containing liquids by connecting h... \n",
"17635 Clean interiors of tank cars or tank trucks, u... \n",
"17636 Lower gauge rods into tanks or read meters to ... \n",
"17637 Operate conveyors and equipment to transfer gr... \n",
"17638 Perform general warehouse activities, such as ... \n",
"\n",
" occupation_title \\\n",
"0 Chief Executives \n",
"1 Chief Executives \n",
"2 Chief Executives \n",
"3 Chief Executives \n",
"4 Chief Executives \n",
"... ... \n",
"17634 Tank Car, Truck, and Ship Loaders \n",
"17635 Tank Car, Truck, and Ship Loaders \n",
"17636 Tank Car, Truck, and Ship Loaders \n",
"17637 Tank Car, Truck, and Ship Loaders \n",
"17638 Tank Car, Truck, and Ship Loaders \n",
"\n",
" occupation_description \\\n",
"0 Determine and formulate policies and provide o... \n",
"1 Determine and formulate policies and provide o... \n",
"2 Determine and formulate policies and provide o... \n",
"3 Determine and formulate policies and provide o... \n",
"4 Determine and formulate policies and provide o... \n",
"... ... \n",
"17634 Load and unload chemicals and bulk solids, suc... \n",
"17635 Load and unload chemicals and bulk solids, suc... \n",
"17636 Load and unload chemicals and bulk solids, suc... \n",
"17637 Load and unload chemicals and bulk solids, suc... \n",
"17638 Load and unload chemicals and bulk solids, suc... \n",
"\n",
" frequency_category_1 frequency_category_2 frequency_category_3 \\\n",
"0 5.92 15.98 29.68 \n",
"1 1.42 14.44 27.31 \n",
"2 15.50 38.21 32.73 \n",
"3 3.03 17.33 20.30 \n",
"4 1.98 14.06 42.60 \n",
"... ... ... ... \n",
"17634 6.05 29.21 6.88 \n",
"17635 1.47 6.33 21.70 \n",
"17636 4.52 1.76 4.65 \n",
"17637 6.97 12.00 2.52 \n",
"17638 5.91 10.85 6.46 \n",
"\n",
" frequency_category_4 frequency_category_5 frequency_category_6 \\\n",
"0 21.18 19.71 4.91 \n",
"1 25.52 26.88 2.52 \n",
"2 5.15 5.25 0.19 \n",
"3 18.10 33.16 2.01 \n",
"4 21.24 13.18 6.24 \n",
"... ... ... ... \n",
"17634 13.95 27.65 7.93 \n",
"17635 25.69 32.35 12.47 \n",
"17636 17.81 37.42 23.31 \n",
"17637 5.90 35.48 22.08 \n",
"17638 14.46 34.14 16.39 \n",
"\n",
" frequency_category_7 importance_average relevance_average OCC_CODE \\\n",
"0 2.63 4.52 74.44 11-1011 \n",
"1 1.90 4.32 81.71 11-1011 \n",
"2 2.98 4.30 93.41 11-1011 \n",
"3 6.07 4.24 97.79 11-1011 \n",
"4 0.70 4.17 92.92 11-1011 \n",
"... ... ... ... ... \n",
"17634 8.34 4.08 64.04 53-7121 \n",
"17635 0.00 4.02 44.33 53-7121 \n",
"17636 10.55 3.88 65.00 53-7121 \n",
"17637 15.05 3.87 47.90 53-7121 \n",
"17638 11.78 3.53 47.84 53-7121 \n",
"\n",
" TOT_EMP H_MEAN A_MEAN \n",
"0 211230.0 124.47 258900 \n",
"1 211230.0 124.47 258900 \n",
"2 211230.0 124.47 258900 \n",
"3 211230.0 124.47 258900 \n",
"4 211230.0 124.47 258900 \n",
"... ... ... ... \n",
"17634 11400.0 29.1 60530 \n",
"17635 11400.0 29.1 60530 \n",
"17636 11400.0 29.1 60530 \n",
"17637 11400.0 29.1 60530 \n",
"17638 11400.0 29.1 60530 \n",
"\n",
"[17639 rows x 18 columns]"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"df_oesm_detailed = df_oesm[df_oesm['O_GROUP'] == 'detailed'][['OCC_CODE', 'TOT_EMP', 'H_MEAN', 'A_MEAN']].copy()\n",
"df_enriched_trs['occ_code_join'] = df_enriched_trs['onetsoc_code'].str[:7]\n",
"df = pd.merge(\n",
" df_enriched_trs,\n",
" df_oesm_detailed,\n",
" left_on='occ_code_join',\n",
" right_on='OCC_CODE',\n",
" how='left'\n",
")\n",
"df = df.drop(columns=['occ_code_join'])\n",
"df"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "9be7acb5-2374-4f61-bba3-13b0077c0bd2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Task: Identify, evaluate and recommend hardware or software technologies to achieve desired database performance.\n",
"Occupation Description: Design strategies for enterprise databases, data warehouse systems, and multidimensional networks. Set standards for database operations, programming, query processes, and security. Model, design, and construct large relational databases or data warehouses. Create and optimize data models for warehouse infrastructure and workflow. Integrate new systems with existing warehouse structure and refine system performance and functionality.\n",
"Occupation Title: Database Architects\n"
]
},
{
"data": {
"text/plain": [
"119976"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"df_merged = pd \\\n",
" .merge(left=df_enriched_trs, right=df_remote_status[['O*NET-SOC Code', 'Remote']], how='left', left_on='onetsoc_code', right_on='O*NET-SOC Code') \\\n",
" .drop(columns=['O*NET-SOC Code']) \\\n",
" .rename(columns={'Remote': 'remote'}) \\\n",
" .rename(columns=FREQUENCY_MAP) \\\n",
" .query('remote == \"remote\" and importance_average >= 3 and relevance_average > 50')\n",
"\n",
"row = df_merged.iloc[30000]\n",
"print('Task: ', row['task'])\n",
"print('Occupation Description: ', row['occupation_description'])\n",
"print('Occupation Title: ', row['occupation_title'])\n",
"\n",
"len(df_merged)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fd9ac1c3-6d17-4764-8a2e-c84d4019bd9e",
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"outputs": [],
"source": [
"# Cross-reference woth BLS OEWS\n",
"# It doesn't really make sens to have it per-task, we only need it per-occupation...\n",
"df_oesm_detailed = df_oesm[df_oesm['O_GROUP'] == 'detailed'][['OCC_CODE', 'TOT_EMP', 'H_MEAN', 'A_MEAN']].copy()\n",
"df_merged['occ_code_join'] = df_merged['onetsoc_code'].str[:7]\n",
"df_merged = pd.merge(\n",
" df_merged,\n",
" df_oesm_detailed,\n",
" left_on='occ_code_join',\n",
" right_on='OCC_CODE',\n",
" how='left'\n",
")\n",
"df_merged = df_merged.drop(columns=['occ_code_join'])\n",
"df_merged"
]
},
{
"cell_type": "code",
"execution_count": 76,
"id": "08f45d91-039d-4ec0-94a2-f305a3312e6a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Why did the scarecrow win an award?\n",
"\n",
"Because he was outstanding in his field!\n"
]
}
],
"source": [
"response = oai.chat.completions.create(messages=[{\"role\": \"user\", \"content\": \"Tell me a joke\"}], model=\"gpt-4.1-2025-04-14\", max_tokens=100, temperature=0.7, n=1, stop=None)\n",
"joke = response.choices[0].message.content.strip()\n",
"print(joke)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

View file

@ -1,24 +0,0 @@
def calc_loss(df):
"""
Geometric-mean log error between prediction bands and golden bands.
Assumes all columns are strictly positive.
Parameters
----------
df : pandas.DataFrame
Must contain the columns:
- 'pred_lower', 'pred_upper'
- 'golden_lower', 'golden_upper'
Returns
-------
float
Scalar loss value (the smaller, the better).
"""
# Element-wise absolute log-ratios
loss_lower = np.abs(np.log(df["pred_lower"] / df["golden_lower"]))
loss_upper = np.abs(np.log(df["pred_upper"] / df["golden_upper"]))
# Average the two means, then exponentiate
loss = np.exp(0.5 * (loss_lower.mean() + loss_upper.mean()))
return loss

View file

@ -1,334 +0,0 @@
import streamlit as st
import sqlite3
import pandas as pd
import graphviz
import textwrap
# --- Database Setup ---
DB_FILE = "onet.database"
@st.cache_resource
def get_db_connection():
"""Establishes a connection to the SQLite database."""
conn = sqlite3.connect(DB_FILE)
conn.row_factory = sqlite3.Row # Access columns by name
return conn
@st.cache_data
def get_occupations(_conn):
"""Fetches all occupations from the database."""
df = pd.read_sql_query(
"SELECT onetsoc_code, title FROM occupation_data ORDER BY title", _conn
)
return df
@st.cache_data
def get_iwas_for_occupation(_conn, onetsoc_code):
"""
Fetches IWAs for a given occupation.
An occupation is linked to Work Activities (element_id in work_activities table).
These Work Activity element_ids are then used in iwa_reference to find associated IWAs.
"""
query = """
SELECT DISTINCT
ir.iwa_id,
ir.iwa_title
FROM work_activities wa
JOIN iwa_reference ir ON wa.element_id = ir.element_id
WHERE wa.onetsoc_code = ?
ORDER BY ir.iwa_title;
"""
df = pd.read_sql_query(query, _conn, params=(onetsoc_code,))
return df
@st.cache_data
def get_dwas_for_iwas(_conn, iwa_ids):
"""Fetches DWAs for a list of IWA IDs."""
if not iwa_ids:
return pd.DataFrame()
placeholders = ",".join(
"?" for _ in iwa_ids
) # Create one placeholder for each IWA ID
query = f"""
SELECT DISTINCT
dr.dwa_id,
dr.dwa_title,
dr.iwa_id -- to link back to the IWA
FROM dwa_reference dr
WHERE dr.iwa_id IN ({placeholders})
ORDER BY dr.dwa_title;
"""
df = pd.read_sql_query(query, _conn, params=iwa_ids)
return df
@st.cache_data
def get_tasks_for_dwas(_conn, onetsoc_code, dwa_ids):
"""Fetches tasks for a given occupation and list of DWA IDs."""
if not dwa_ids:
return pd.DataFrame()
placeholders = ",".join(
"?" for _ in dwa_ids
) # Create one placeholder for each DWA ID
query = f"""
SELECT DISTINCT
ts.task_id,
ts.task,
t2d.dwa_id -- to link back to the DWA
FROM task_statements ts
JOIN tasks_to_dwas t2d ON ts.task_id = t2d.task_id
WHERE ts.onetsoc_code = ? AND t2d.dwa_id IN ({placeholders})
ORDER BY ts.task;
"""
# The parameters list should first contain onetsoc_code, then all DWA IDs.
params = [onetsoc_code] + dwa_ids
df = pd.read_sql_query(query, _conn, params=params)
return df
def smart_wrap(text, width=40):
"""Wraps text for better display in graph nodes."""
return "\n".join(
textwrap.wrap(
text,
width=width,
break_long_words=True,
replace_whitespace=False,
drop_whitespace=False,
)
)
# --- Streamlit App Layout ---
st.set_page_config(layout="wide")
# Check if database file exists
try:
# Attempt to open for binary read to check existence and basic readability
with open(DB_FILE, "rb") as f:
pass
conn = get_db_connection()
except FileNotFoundError:
st.error(
f"Database file '{DB_FILE}' not found. Please ensure it is in the same directory as the script."
)
st.stop()
except sqlite3.Error as e:
st.error(f"Error connecting to or reading the database '{DB_FILE}': {e}")
st.info(
"Please ensure the database file is a valid SQLite database and not corrupted."
)
st.stop()
st.title("O*NET Occupation Hierarchy Explorer")
st.markdown("""
This application visualizes the relationships between Occupations, Intermediate Work Activities (IWAs),
Detailed Work Activities (DWAs), and Task Statements from the O*NET database.
Select an occupation from the control panel on the left to view its hierarchical breakdown.
""")
# --- Sidebar for Occupation Selection ---
col1, col2 = st.columns([0.3, 0.7], gap="large")
with col1:
st.header("Control Panel")
occupations_df = get_occupations(conn)
if occupations_df.empty:
st.warning("No occupations found in the database.")
st.stop()
# Create a display string with code and title for the selectbox
occupations_df["display_name"] = (
occupations_df["title"] + " (" + occupations_df["onetsoc_code"] + ")"
)
search_term = st.text_input(
"Search for an occupation:", placeholder="E.g., Software Developer"
)
if search_term:
# Ensure search term is treated as a literal string for regex, if needed, or use basic string methods
search_term_safe = (
search_term.replace("[", "\\[")
.replace("]", "\\]")
.replace("(", "\\(")
.replace(")", "\\)")
)
filtered_occupations = occupations_df[
occupations_df["title"].str.contains(
search_term_safe, case=False, regex=True
)
| occupations_df["onetsoc_code"].str.contains(
search_term_safe, case=False, regex=True
)
]
else:
filtered_occupations = occupations_df
if not filtered_occupations.empty:
# Sort filtered occupations for consistent display in selectbox
filtered_occupations_sorted = filtered_occupations.sort_values("display_name")
selected_occupation_display_name = st.selectbox(
"Choose an occupation:",
options=filtered_occupations_sorted["display_name"],
index=0, # Default to the first item
)
# Get the onetsoc_code and title from the selected display name
selected_row = occupations_df[
occupations_df["display_name"] == selected_occupation_display_name
].iloc[0]
selected_onetsoc_code = selected_row["onetsoc_code"]
selected_occupation_title = selected_row["title"]
else:
st.warning("No occupations match your search term.")
selected_onetsoc_code = None
selected_occupation_title = None
# --- Main Area for Graph Display ---
with col2:
st.header("Occupation Graph")
if selected_onetsoc_code:
st.subheader(
f"Displaying: {selected_occupation_title} ({selected_onetsoc_code})"
)
iwas_df = get_iwas_for_occupation(conn, selected_onetsoc_code)
if iwas_df.empty:
st.info(
"No Intermediate Work Activities (IWAs) found directly linked for this occupation."
)
else:
graph = graphviz.Digraph(
comment=f"O*NET Hierarchy for {selected_onetsoc_code}"
)
graph.attr(
rankdir="LR",
splines="spline",
concentrate="false",
nodesep="0.5",
ranksep="0.8",
)
# Occupation Node
occ_node_id = f"occ_{selected_onetsoc_code.replace('.', '_')}" # Ensure ID is valid for DOT
occ_label = smart_wrap(
f"Occupation: {selected_occupation_title}\n({selected_onetsoc_code})",
width=30,
)
graph.node(
occ_node_id,
label=occ_label,
shape="ellipse",
style="filled",
fillcolor="skyblue",
)
# Fetch DWAs
iwa_ids = iwas_df["iwa_id"].tolist()
dwas_df = get_dwas_for_iwas(conn, iwa_ids)
dwa_ids_for_tasks = []
if not dwas_df.empty:
dwa_ids_for_tasks = dwas_df["dwa_id"].unique().tolist()
# Fetch Tasks
tasks_df = get_tasks_for_dwas(
conn, selected_onetsoc_code, dwa_ids_for_tasks
)
# Add IWA Nodes and Edges
for _, iwa_row in iwas_df.iterrows():
iwa_node_id = f"iwa_{str(iwa_row['iwa_id']).replace('.', '_')}"
iwa_label = smart_wrap(
f"IWA: {iwa_row['iwa_title']}\n(ID: {iwa_row['iwa_id']})", width=35
)
graph.node(
iwa_node_id,
label=iwa_label,
shape="box",
style="filled",
fillcolor="khaki",
)
graph.edge(occ_node_id, iwa_node_id)
# Add DWA Nodes and Edges (for this IWA)
current_iwa_dwas = dwas_df[dwas_df["iwa_id"] == iwa_row["iwa_id"]]
for _, dwa_row in current_iwa_dwas.iterrows():
dwa_node_id = f"dwa_{str(dwa_row['dwa_id']).replace('.', '_')}"
dwa_label = smart_wrap(
f"DWA: {dwa_row['dwa_title']}\n(ID: {dwa_row['dwa_id']})",
width=40,
)
graph.node(
dwa_node_id,
label=dwa_label,
shape="box",
style="filled",
fillcolor="lightcoral",
)
graph.edge(iwa_node_id, dwa_node_id)
# Add Task Nodes and Edges (for this DWA and Occupation)
current_dwa_tasks = tasks_df[
tasks_df["dwa_id"] == dwa_row["dwa_id"]
]
for _, task_row in current_dwa_tasks.iterrows():
# Ensure task_id is a string and valid for DOT
task_id_str = str(task_row["task_id"]).split(".")[
0
] # Handle decimal task_ids if they appear
task_node_id = f"task_{task_id_str}"
task_label = smart_wrap(
f"Task: {task_row['task']}\n(ID: {task_id_str})", width=50
)
graph.node(
task_node_id,
label=task_label,
shape="note",
style="filled",
fillcolor="lightgray",
)
graph.edge(dwa_node_id, task_node_id)
if (
not graph.body or len(graph.body) <= 1
): # Check if any nodes were actually added beyond the occupation
st.info(
"No hierarchical data (IWAs, DWAs, Tasks) to display for this occupation after initial selection."
)
else:
try:
st.graphviz_chart(graph, use_container_width=True)
with st.expander("View Data Tables for Selected Occupation"):
st.markdown("##### Intermediate Work Activities (IWAs)")
st.dataframe(iwas_df, use_container_width=True)
if not dwas_df.empty:
st.markdown("##### Detailed Work Activities (DWAs)")
st.dataframe(dwas_df, use_container_width=True)
if not tasks_df.empty:
st.markdown("##### Task Statements")
st.dataframe(tasks_df, use_container_width=True)
except Exception as e:
st.error(
f"Could not render the graph. Graphviz might not be installed correctly or there's an issue with the graph data: {e}"
)
st.text("Graphviz DOT source (for debugging):")
st.code(graph.source, language="dot")
else:
st.info("Select an occupation from the control panel to see its graph.")
# Instructions to run the app:
# 1. Save this code as a Python file (e.g., onet_explorer_app.py).
# 2. Ensure the 'onet.database' file is in the same directory.
# 3. Install the required libraries: pip install streamlit pandas graphviz
# 4. Open your terminal or command prompt, navigate to the directory, and run:
# streamlit run onet_explorer_app.py

View file

@ -1,352 +0,0 @@
CREATE TABLE content_model_reference (
element_id CHARACTER VARYING(20) NOT NULL,
element_name CHARACTER VARYING(150) NOT NULL,
description CHARACTER VARYING(1500) NOT NULL,
PRIMARY KEY (element_id));
CREATE TABLE job_zone_reference (
job_zone DECIMAL(1,0) NOT NULL,
name CHARACTER VARYING(50) NOT NULL,
experience CHARACTER VARYING(300) NOT NULL,
education CHARACTER VARYING(500) NOT NULL,
job_training CHARACTER VARYING(300) NOT NULL,
examples CHARACTER VARYING(500) NOT NULL,
svp_range CHARACTER VARYING(25) NOT NULL,
PRIMARY KEY (job_zone));
CREATE TABLE occupation_data (
onetsoc_code CHARACTER(10) NOT NULL,
title CHARACTER VARYING(150) NOT NULL,
description CHARACTER VARYING(1000) NOT NULL,
PRIMARY KEY (onetsoc_code));
CREATE TABLE scales_reference (
scale_id CHARACTER VARYING(3) NOT NULL,
scale_name CHARACTER VARYING(50) NOT NULL,
minimum DECIMAL(1,0) NOT NULL,
maximum DECIMAL(3,0) NOT NULL,
PRIMARY KEY (scale_id));
CREATE TABLE ete_categories (
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
category DECIMAL(3,0) NOT NULL,
category_description CHARACTER VARYING(1000) NOT NULL,
PRIMARY KEY (element_id, scale_id, category),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE level_scale_anchors (
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
anchor_value DECIMAL(3,0) NOT NULL,
anchor_description CHARACTER VARYING(1000) NOT NULL,
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE occupation_level_metadata (
onetsoc_code CHARACTER(10) NOT NULL,
item CHARACTER VARYING(150) NOT NULL,
response CHARACTER VARYING(75),
n DECIMAL(4,0),
percent DECIMAL(4,1),
date_updated DATE NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code));
CREATE TABLE survey_booklet_locations (
element_id CHARACTER VARYING(20) NOT NULL,
survey_item_number CHARACTER VARYING(4) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE task_categories (
scale_id CHARACTER VARYING(3) NOT NULL,
category DECIMAL(3,0) NOT NULL,
category_description CHARACTER VARYING(1000) NOT NULL,
PRIMARY KEY (scale_id, category),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE work_context_categories (
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
category DECIMAL(3,0) NOT NULL,
category_description CHARACTER VARYING(1000) NOT NULL,
PRIMARY KEY (element_id, scale_id, category),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE abilities (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
not_relevant CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE education_training_experience (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
category DECIMAL(3,0),
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id),
FOREIGN KEY (element_id, scale_id, category) REFERENCES ete_categories(element_id, scale_id, category));
CREATE TABLE interests (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
data_value DECIMAL(5,2) NOT NULL,
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE job_zones (
onetsoc_code CHARACTER(10) NOT NULL,
job_zone DECIMAL(1,0) NOT NULL,
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (job_zone) REFERENCES job_zone_reference(job_zone));
CREATE TABLE knowledge (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
not_relevant CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE skills (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
not_relevant CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE task_statements (
onetsoc_code CHARACTER(10) NOT NULL,
task_id DECIMAL(8,0) NOT NULL,
task CHARACTER VARYING(1000) NOT NULL,
task_type CHARACTER VARYING(12),
incumbents_responding DECIMAL(4,0),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
PRIMARY KEY (task_id),
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code));
CREATE TABLE task_ratings (
onetsoc_code CHARACTER(10) NOT NULL,
task_id DECIMAL(8,0) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
category DECIMAL(3,0),
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (task_id) REFERENCES task_statements(task_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id),
FOREIGN KEY (scale_id, category) REFERENCES task_categories(scale_id, category));
CREATE TABLE work_activities (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
not_relevant CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE work_context (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
category DECIMAL(3,0),
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
not_relevant CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id),
FOREIGN KEY (element_id, scale_id, category) REFERENCES work_context_categories(element_id, scale_id, category));
CREATE TABLE work_styles (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
data_value DECIMAL(5,2) NOT NULL,
n DECIMAL(4,0),
standard_error DECIMAL(7,4),
lower_ci_bound DECIMAL(7,4),
upper_ci_bound DECIMAL(7,4),
recommend_suppress CHARACTER(1),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE work_values (
onetsoc_code CHARACTER(10) NOT NULL,
element_id CHARACTER VARYING(20) NOT NULL,
scale_id CHARACTER VARYING(3) NOT NULL,
data_value DECIMAL(5,2) NOT NULL,
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (scale_id) REFERENCES scales_reference(scale_id));
CREATE TABLE iwa_reference (
element_id CHARACTER VARYING(20) NOT NULL,
iwa_id CHARACTER VARYING(20) NOT NULL,
iwa_title CHARACTER VARYING(150) NOT NULL,
PRIMARY KEY (iwa_id),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE dwa_reference (
element_id CHARACTER VARYING(20) NOT NULL,
iwa_id CHARACTER VARYING(20) NOT NULL,
dwa_id CHARACTER VARYING(20) NOT NULL,
dwa_title CHARACTER VARYING(150) NOT NULL,
PRIMARY KEY (dwa_id),
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (iwa_id) REFERENCES iwa_reference(iwa_id));
CREATE TABLE tasks_to_dwas (
onetsoc_code CHARACTER(10) NOT NULL,
task_id DECIMAL(8,0) NOT NULL,
dwa_id CHARACTER VARYING(20) NOT NULL,
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (task_id) REFERENCES task_statements(task_id),
FOREIGN KEY (dwa_id) REFERENCES dwa_reference(dwa_id));
CREATE TABLE emerging_tasks (
onetsoc_code CHARACTER(10) NOT NULL,
task CHARACTER VARYING(1000) NOT NULL,
category CHARACTER VARYING(8) NOT NULL,
original_task_id DECIMAL(8,0),
date_updated DATE NOT NULL,
domain_source CHARACTER VARYING(30) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (original_task_id) REFERENCES task_statements(task_id));
CREATE TABLE related_occupations (
onetsoc_code CHARACTER(10) NOT NULL,
related_onetsoc_code CHARACTER(10) NOT NULL,
relatedness_tier CHARACTER VARYING(50) NOT NULL,
related_index DECIMAL(3,0) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (related_onetsoc_code) REFERENCES occupation_data(onetsoc_code));
CREATE TABLE unspsc_reference (
commodity_code DECIMAL(8,0) NOT NULL,
commodity_title CHARACTER VARYING(150) NOT NULL,
class_code DECIMAL(8,0) NOT NULL,
class_title CHARACTER VARYING(150) NOT NULL,
family_code DECIMAL(8,0) NOT NULL,
family_title CHARACTER VARYING(150) NOT NULL,
segment_code DECIMAL(8,0) NOT NULL,
segment_title CHARACTER VARYING(150) NOT NULL,
PRIMARY KEY (commodity_code));
CREATE TABLE alternate_titles (
onetsoc_code CHARACTER(10) NOT NULL,
alternate_title CHARACTER VARYING(250) NOT NULL,
short_title CHARACTER VARYING(150),
sources CHARACTER VARYING(50) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code));
CREATE TABLE sample_of_reported_titles (
onetsoc_code CHARACTER(10) NOT NULL,
reported_job_title CHARACTER VARYING(150) NOT NULL,
shown_in_my_next_move CHARACTER(1) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code));
CREATE TABLE technology_skills (
onetsoc_code CHARACTER(10) NOT NULL,
example CHARACTER VARYING(150) NOT NULL,
commodity_code DECIMAL(8,0) NOT NULL,
hot_technology CHARACTER(1) NOT NULL,
in_demand CHARACTER(1) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (commodity_code) REFERENCES unspsc_reference(commodity_code));
CREATE TABLE tools_used (
onetsoc_code CHARACTER(10) NOT NULL,
example CHARACTER VARYING(150) NOT NULL,
commodity_code DECIMAL(8,0) NOT NULL,
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code),
FOREIGN KEY (commodity_code) REFERENCES unspsc_reference(commodity_code));
CREATE TABLE abilities_to_work_activities (
abilities_element_id CHARACTER VARYING(20) NOT NULL,
work_activities_element_id CHARACTER VARYING(20) NOT NULL,
FOREIGN KEY (abilities_element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (work_activities_element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE abilities_to_work_context (
abilities_element_id CHARACTER VARYING(20) NOT NULL,
work_context_element_id CHARACTER VARYING(20) NOT NULL,
FOREIGN KEY (abilities_element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (work_context_element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE skills_to_work_activities (
skills_element_id CHARACTER VARYING(20) NOT NULL,
work_activities_element_id CHARACTER VARYING(20) NOT NULL,
FOREIGN KEY (skills_element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (work_activities_element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE skills_to_work_context (
skills_element_id CHARACTER VARYING(20) NOT NULL,
work_context_element_id CHARACTER VARYING(20) NOT NULL,
FOREIGN KEY (skills_element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (work_context_element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE riasec_keywords (
element_id CHARACTER VARYING(20) NOT NULL,
keyword CHARACTER VARYING(150) NOT NULL,
keyword_type CHARACTER VARYING(20) NOT NULL,
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE basic_interests_to_riasec (
basic_interests_element_id CHARACTER VARYING(20) NOT NULL,
riasec_element_id CHARACTER VARYING(20) NOT NULL,
FOREIGN KEY (basic_interests_element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (riasec_element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE interests_illus_activities (
element_id CHARACTER VARYING(20) NOT NULL,
interest_type CHARACTER VARYING(20) NOT NULL,
activity CHARACTER VARYING(150) NOT NULL,
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id));
CREATE TABLE interests_illus_occupations (
element_id CHARACTER VARYING(20) NOT NULL,
interest_type CHARACTER VARYING(20) NOT NULL,
onetsoc_code CHARACTER(10) NOT NULL,
FOREIGN KEY (element_id) REFERENCES content_model_reference(element_id),
FOREIGN KEY (onetsoc_code) REFERENCES occupation_data(onetsoc_code));
CREATE TABLE sqlite_stat1(tbl,idx,stat);

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -1,411 +0,0 @@
import pandas as pd
import litellm
import dotenv
import os
import time
import json
import math
# Load environment variables
dotenv.load_dotenv(override=True)
# litellm._turn_on_debug() # Optional debugging
# --- Configuration ---
MODEL = "gpt-4.1-mini" # Make sure this model supports json_schema or structured output
RATE_LIMIT = 5000 # Requests per minute
CHUNK_SIZE = 300 # Number of unique tasks per API call
SECONDS_PER_MINUTE = 60
# File configuration
CLASSIFICATION_FILENAME = "tasks_estimateable.csv" # Output file with classifications
TASK_SOURCE_FOR_INIT_FILENAME = "tasks_with_estimates.csv"
OUTPUT_COLUMN_NAME = "task_estimateable"
SOURCE_FILTER_COLUMN = "remote_status"
SOURCE_FILTER_VALUE = "remote"
# --- Prompts and Schema ---
SYSTEM_PROMPT_CLASSIFY = """
Classify the provided O*NET task into one of these categories:
- ATOMIC (schedulable): A single, clearly-bounded activity, typically lasting minutes, hours, or a few days.
- ONGOING-CONSTRAINT (background role/ethical rule): A continuous responsibility or behavioural norm with no schedulable duration (e.g., follow confidentiality rules, serve as department head).
""".strip()
USER_MESSAGE_TEMPLATE_CLASSIFY = "Task: {task}"
CLASSIFICATION_CATEGORIES = ["ATOMIC", "ONGOING-CONSTRAINT"]
SCHEMA_FOR_CLASSIFICATION = {
"name": "classify_task_type",
"strict": True,
"schema": {
"type": "object",
"properties": {
"task_category": {
"type": "string",
"enum": CLASSIFICATION_CATEGORIES,
"description": "The classification of the task (ATOMIC or ONGOING-CONSTRAINT).",
}
},
"required": ["task_category"],
"additionalProperties": False,
},
}
def save_dataframe(df_to_save, filename):
"""Saves the DataFrame to the specified CSV file using atomic write."""
try:
temp_filename = filename + ".tmp"
df_to_save.to_csv(temp_filename, encoding="utf-8-sig", index=False)
os.replace(temp_filename, filename)
except Exception as e:
print(f"--- Error saving DataFrame to {filename}: {e} ---")
if os.path.exists(temp_filename):
try:
os.remove(temp_filename)
except Exception as remove_err:
print(
f"--- Error removing temporary save file {temp_filename}: {remove_err} ---"
)
# --- Load or Initialize DataFrame ---
try:
if os.path.exists(CLASSIFICATION_FILENAME):
df = pd.read_csv(CLASSIFICATION_FILENAME, encoding="utf-8-sig")
print(f"Successfully read {len(df)} rows from {CLASSIFICATION_FILENAME}.")
save_needed_after_load = False
if OUTPUT_COLUMN_NAME not in df.columns:
df[OUTPUT_COLUMN_NAME] = pd.NA
print(f"Added '{OUTPUT_COLUMN_NAME}' column.")
save_needed_after_load = True
df[OUTPUT_COLUMN_NAME].replace(["", None, ""], pd.NA, inplace=True)
if df[OUTPUT_COLUMN_NAME].dtype != object and not isinstance(
df[OUTPUT_COLUMN_NAME].dtype, pd.StringDtype
):
try:
df[OUTPUT_COLUMN_NAME] = df[OUTPUT_COLUMN_NAME].astype(object)
print(
f"Corrected dtype of '{OUTPUT_COLUMN_NAME}' to {df[OUTPUT_COLUMN_NAME].dtype}."
)
save_needed_after_load = True
except Exception as e:
print(
f"Warning: Could not convert column '{OUTPUT_COLUMN_NAME}' to object: {e}."
)
if "task" not in df.columns:
print(
f"Error: {CLASSIFICATION_FILENAME} must contain a 'task' column for processing."
)
exit()
if save_needed_after_load:
print(f"Saving {CLASSIFICATION_FILENAME} after adding/adjusting column.")
save_dataframe(df, CLASSIFICATION_FILENAME)
else:
print(
f"{CLASSIFICATION_FILENAME} not found. Attempting to create it from {TASK_SOURCE_FOR_INIT_FILENAME}."
)
if not os.path.exists(TASK_SOURCE_FOR_INIT_FILENAME):
print(
f"Error: Source file {TASK_SOURCE_FOR_INIT_FILENAME} not found. Cannot create {CLASSIFICATION_FILENAME}."
)
exit()
df_source = pd.read_csv(TASK_SOURCE_FOR_INIT_FILENAME, encoding="utf-8-sig")
required_source_cols_for_init = ["task", SOURCE_FILTER_COLUMN]
missing_source_cols = [
col for col in required_source_cols_for_init if col not in df_source.columns
]
if missing_source_cols:
print(
f"Error: Source file {TASK_SOURCE_FOR_INIT_FILENAME} is missing required columns for initialization: {', '.join(missing_source_cols)}."
)
exit()
df_source_filtered = df_source[
df_source[SOURCE_FILTER_COLUMN] == SOURCE_FILTER_VALUE
].copy()
if df_source_filtered.empty:
print(
f"Warning: No tasks with '{SOURCE_FILTER_COLUMN}' == '{SOURCE_FILTER_VALUE}' found in {TASK_SOURCE_FOR_INIT_FILENAME}. "
f"{CLASSIFICATION_FILENAME} will be created with schema but no tasks to classify initially."
)
df = df_source_filtered[["task"]].copy()
df[OUTPUT_COLUMN_NAME] = pd.NA
df[OUTPUT_COLUMN_NAME] = df[OUTPUT_COLUMN_NAME].astype(object)
print(
f"Created {CLASSIFICATION_FILENAME} using tasks from {TASK_SOURCE_FOR_INIT_FILENAME} "
f"(where {SOURCE_FILTER_COLUMN}='{SOURCE_FILTER_VALUE}'). New file has {len(df)} tasks."
)
save_dataframe(df, CLASSIFICATION_FILENAME)
except FileNotFoundError:
print(f"Error: A required file was not found. Please check paths.")
exit()
except Exception as e:
print(f"Error during DataFrame loading or initialization: {e}")
exit()
# --- Identify Unique Tasks to Process ---
if df.empty:
print(f"{CLASSIFICATION_FILENAME} is empty. Nothing to process. Exiting.")
exit()
initial_unprocessed_mask = df[OUTPUT_COLUMN_NAME].isna()
if not initial_unprocessed_mask.any():
print(
f"All tasks in {CLASSIFICATION_FILENAME} seem to have been classified already. Exiting."
)
exit()
# Filter for rows that are unprocessed AND have a valid 'task' string
valid_tasks_to_consider_df = df[
initial_unprocessed_mask & df["task"].notna() & (df["task"].str.strip() != "")
]
if valid_tasks_to_consider_df.empty:
print(
f"No valid, unclassified tasks found to process (after filtering out empty/NaN task descriptions). Exiting."
)
exit()
unique_task_labels_for_api = (
valid_tasks_to_consider_df["task"].drop_duplicates().tolist()
)
total_rows_to_update_potentially = len(
df[initial_unprocessed_mask]
) # Count all rows that are NA
print(
f"Found {total_rows_to_update_potentially} total rows in {CLASSIFICATION_FILENAME} needing classification."
)
print(
f"Identified {len(unique_task_labels_for_api)} unique, valid task labels to send to the API."
)
# --- Prepare messages for batch completion (only for unique task labels) ---
messages_list = []
print(f"Preparing messages for {len(unique_task_labels_for_api)} unique task labels...")
for task_label in unique_task_labels_for_api:
# task_label is already guaranteed to be non-empty and not NaN from the filtering above
user_message = USER_MESSAGE_TEMPLATE_CLASSIFY.format(task=task_label)
messages_for_task = [
{"role": "system", "content": SYSTEM_PROMPT_CLASSIFY},
{"role": "user", "content": user_message},
]
messages_list.append(messages_for_task)
print(f"Prepared {len(messages_list)} message sets for batch completion.")
if (
not messages_list
): # Should only happen if unique_task_labels_for_api was empty, caught above
print(
"No messages prepared, though unique tasks were identified. This is unexpected. Exiting."
)
exit()
# --- Call batch_completion in chunks with rate limiting and periodic saving ---
total_unique_tasks_to_send = len(
messages_list
) # Same as len(unique_task_labels_for_api)
num_chunks = math.ceil(total_unique_tasks_to_send / CHUNK_SIZE)
print(
f"\nStarting batch classification for {total_unique_tasks_to_send} unique task labels in {num_chunks} chunks..."
)
overall_start_time = time.time()
processed_rows_count_total = 0 # Counts actual rows updated in the DataFrame
for i in range(num_chunks):
chunk_start_message_index = i * CHUNK_SIZE
chunk_end_message_index = min((i + 1) * CHUNK_SIZE, total_unique_tasks_to_send)
message_chunk = messages_list[chunk_start_message_index:chunk_end_message_index]
# Get corresponding unique task labels for this chunk
chunk_task_labels = unique_task_labels_for_api[
chunk_start_message_index:chunk_end_message_index
]
if not message_chunk: # Should not happen if loop range is correct
continue
print(
f"\nProcessing chunk {i + 1}/{num_chunks} (Unique Task Labels {chunk_start_message_index + 1}-{chunk_end_message_index} of this run)..."
)
chunk_start_time = time.time()
responses = []
try:
print(
f"Sending {len(message_chunk)} requests (for unique tasks) for chunk {i + 1}..."
)
responses = litellm.batch_completion(
model=MODEL,
messages=message_chunk,
response_format={
"type": "json_schema",
"json_schema": SCHEMA_FOR_CLASSIFICATION,
},
num_retries=3,
)
print(f"Chunk {i + 1} API call completed.")
except Exception as e:
print(f"Error during litellm.batch_completion for chunk {i + 1}: {e}")
responses = [None] * len(message_chunk)
# --- Process responses for the current chunk ---
# chunk_updates stores {task_label: classification_category}
chunk_task_classifications = {}
successful_api_calls_in_chunk = 0
failed_api_calls_in_chunk = 0
if responses and len(responses) == len(message_chunk):
for j, response in enumerate(responses):
current_task_label = chunk_task_labels[
j
] # The unique task label for this response
content_str = None
if response is None:
print(
f"API call failed for task label '{current_task_label}' (response is None)."
)
failed_api_calls_in_chunk += 1
continue
try:
if (
response.choices
and response.choices[0].message
and response.choices[0].message.content
):
content_str = response.choices[0].message.content
classification_data = json.loads(content_str)
category_raw = classification_data.get("task_category")
if category_raw in CLASSIFICATION_CATEGORIES:
successful_api_calls_in_chunk += 1
chunk_task_classifications[current_task_label] = category_raw
else:
print(
f"Warning: Invalid or missing task_category for task label '{current_task_label}': '{category_raw}'. Content: '{content_str}'"
)
failed_api_calls_in_chunk += 1
else:
finish_reason = (
response.choices[0].finish_reason
if (response.choices and response.choices[0].finish_reason)
else "unknown"
)
error_message = (
response.choices[0].message.content
if (response.choices and response.choices[0].message)
else "No content in message."
)
print(
f"Warning: Received non-standard or empty response content for task label '{current_task_label}'. "
f"Finish Reason: '{finish_reason}'. Message: '{error_message}'. Raw Choices: {response.choices}"
)
failed_api_calls_in_chunk += 1
except json.JSONDecodeError:
print(
f"Warning: Could not decode JSON for task label '{current_task_label}'. Content received: '{content_str}'"
)
failed_api_calls_in_chunk += 1
except AttributeError as ae:
print(
f"Warning: Missing attribute processing response for task label '{current_task_label}': {ae}. Response: {response}"
)
failed_api_calls_in_chunk += 1
except Exception as e:
print(
f"Warning: Unexpected error processing response for task label '{current_task_label}': {type(e).__name__} - {e}. Response: {response}"
)
failed_api_calls_in_chunk += 1
else:
print(
f"Warning: Mismatch between #responses ({len(responses) if responses else 0}) "
f"and #messages sent ({len(message_chunk)}) for chunk {i + 1}, or no responses. Marking all API calls in chunk as failed."
)
failed_api_calls_in_chunk = len(message_chunk)
# --- Update Main DataFrame and Save Periodically ---
rows_updated_this_chunk = 0
if chunk_task_classifications:
print(
f"Updating main DataFrame with classifications for {len(chunk_task_classifications)} unique tasks from chunk {i + 1}..."
)
for task_label, category in chunk_task_classifications.items():
# Update all rows in the main df that match this task_label AND are still NA in the output column
update_condition = (df["task"] == task_label) & (
df[OUTPUT_COLUMN_NAME].isna()
)
num_rows_for_this_task_label = df[update_condition].shape[0]
if num_rows_for_this_task_label > 0:
df.loc[update_condition, OUTPUT_COLUMN_NAME] = category
rows_updated_this_chunk += num_rows_for_this_task_label
print(
f"Updated {rows_updated_this_chunk} rows in the DataFrame based on this chunk's API responses."
)
print(f"Saving progress to {CLASSIFICATION_FILENAME}...")
save_dataframe(df, CLASSIFICATION_FILENAME)
else:
print(
f"No successful API classifications obtained in chunk {i + 1} to update DataFrame or save."
)
print(
f"Chunk {i + 1} API summary: Successful Calls={successful_api_calls_in_chunk}, Failed/Skipped Calls={failed_api_calls_in_chunk}. "
f"Rows updated in DataFrame this chunk: {rows_updated_this_chunk}"
)
processed_rows_count_total += rows_updated_this_chunk
# --- Rate Limiting Pause ---
chunk_end_time = time.time()
chunk_duration = chunk_end_time - chunk_start_time
print(f"Chunk {i + 1} (API calls and DF update) took {chunk_duration:.2f} seconds.")
if i < num_chunks - 1:
time_per_request = SECONDS_PER_MINUTE / RATE_LIMIT if RATE_LIMIT > 0 else 0
min_chunk_duration_for_rate = (
len(message_chunk) * time_per_request
) # Based on API calls made
pause_needed = max(0, min_chunk_duration_for_rate - chunk_duration)
if pause_needed > 0:
print(
f"Pausing for {pause_needed:.2f} seconds to respect rate limit ({RATE_LIMIT}/min)..."
)
time.sleep(pause_needed)
overall_end_time = time.time()
total_duration_minutes = (overall_end_time - overall_start_time) / 60
print(
f"\nBatch classification finished."
f" Updated {processed_rows_count_total} rows in '{CLASSIFICATION_FILENAME}' with new classifications in this run."
f" Total duration: {total_duration_minutes:.2f} minutes."
)
print(f"Performing final save to {CLASSIFICATION_FILENAME}...")
save_dataframe(df, CLASSIFICATION_FILENAME)
print("\nScript finished.")

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

Binary file not shown.

BIN
dist/estimate_distribution_histplot.png vendored Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 75 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 295 KiB

BIN
dist/estimates_spread_per_occupation.png vendored Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 279 KiB

BIN
dist/intermediate/df_tasks.parquet vendored Normal file

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

After

Width:  |  Height:  |  Size: 239 KiB

BIN
dist/projected_task_automation_p50.png vendored Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 145 KiB

BIN
dist/projected_task_automation_p80.png vendored Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 136 KiB

BIN
dist/sequential_coherence_cdf.png vendored Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 145 KiB

View file

@ -1,392 +0,0 @@
import sqlite3
import pandas as pd
import json
import os
from collections import defaultdict
import numpy as np
# --- Configuration ---
DB_FILE = "onet.database"
OUTPUT_FILE = "task_ratings_enriched.json" # Changed output filename
# --- Database Interaction ---
def fetch_data_from_db(db_path):
"""
Fetches required data from the O*NET SQLite database using JOINs,
including DWAs.
Args:
db_path (str): Path to the SQLite database file.
Returns:
tuple(pandas.DataFrame, pandas.DataFrame): A tuple containing:
- DataFrame with task ratings info.
- DataFrame with task-to-DWA mapping.
Returns (None, None) if the database file doesn't exist or an error occurs.
"""
if not os.path.exists(db_path):
print(f"Error: Database file not found at {db_path}")
return None, None
try:
conn = sqlite3.connect(db_path)
# Construct the SQL query to join the tables and select necessary columns
# Added LEFT JOINs for tasks_to_dwas and dwa_reference
# Use LEFT JOIN in case a task has no DWAs
query = """
SELECT
tr.onetsoc_code,
tr.task_id,
ts.task,
od.title AS occupation_title,
od.description AS occupation_description,
tr.scale_id,
tr.category,
tr.data_value,
dr.dwa_title -- Added DWA title
FROM
task_ratings tr
JOIN
task_statements ts ON tr.task_id = ts.task_id
JOIN
occupation_data od ON tr.onetsoc_code = od.onetsoc_code
LEFT JOIN
tasks_to_dwas td ON tr.onetsoc_code = td.onetsoc_code AND tr.task_id = td.task_id --
LEFT JOIN
dwa_reference dr ON td.dwa_id = dr.dwa_id; --
"""
df = pd.read_sql_query(query, conn)
conn.close()
print(
f"Successfully fetched {len(df)} records (including DWA info) from the database."
)
if df.empty:
print("Warning: Fetched DataFrame is empty.")
# Return empty DataFrames with expected columns if the main fetch is empty
ratings_cols = [
"onetsoc_code",
"task_id",
"task",
"occupation_title",
"occupation_description",
"scale_id",
"category",
"data_value",
]
dwa_cols = ["onetsoc_code", "task_id", "dwa_title"]
return pd.DataFrame(columns=ratings_cols), pd.DataFrame(columns=dwa_cols)
# Remove duplicates caused by joining ratings with potentially multiple DWAs per task
# Keep only unique combinations of the core task/rating info before processing
core_cols = [
"onetsoc_code",
"task_id",
"task",
"occupation_title",
"occupation_description",
"scale_id",
"category",
"data_value",
]
# Check if all core columns exist before attempting to drop duplicates
missing_core_cols = [col for col in core_cols if col not in df.columns]
if missing_core_cols:
print(f"Error: Missing core columns in fetched data: {missing_core_cols}")
return None, None
ratings_df = df[core_cols].drop_duplicates().reset_index(drop=True)
# Get unique DWA info separately
dwa_cols = ["onetsoc_code", "task_id", "dwa_title"]
# Check if all DWA columns exist before processing
if all(col in df.columns for col in dwa_cols):
dwas_df = (
df[dwa_cols]
.dropna(subset=["dwa_title"])
.drop_duplicates()
.reset_index(drop=True)
)
else:
print("Warning: DWA related columns missing, creating empty DWA DataFrame.")
dwas_df = pd.DataFrame(
columns=dwa_cols
) # Create empty df if columns missing
return ratings_df, dwas_df # Return two dataframes now
except sqlite3.Error as e:
print(f"SQLite error: {e}")
if "conn" in locals() and conn:
conn.close()
return None, None # Return None for both if error
except Exception as e:
print(f"An error occurred during data fetching: {e}")
if "conn" in locals() and conn:
conn.close()
return None, None # Return None for both if error
# --- Data Processing ---
def process_task_ratings_with_dwas(ratings_df, dwas_df):
"""
Processes the fetched data to group, pivot frequency, calculate averages,
structure the output, and add associated DWAs.
Args:
ratings_df (pandas.DataFrame): The input DataFrame with task ratings info.
dwas_df (pandas.DataFrame): The input DataFrame with task-to-DWA mapping. Can be None or empty.
Returns:
list: A list of dictionaries, each representing an enriched task rating with DWAs.
Returns None if the input ratings DataFrame is invalid.
"""
if ratings_df is None or not isinstance(
ratings_df, pd.DataFrame
): # Check if it's a DataFrame
print("Error: Input ratings DataFrame is invalid.")
return None
if ratings_df.empty:
print(
"Warning: Input ratings DataFrame is empty. Processing will yield empty result."
)
# Decide how to handle empty input, maybe return empty list directly
# return []
# Ensure dwas_df is a DataFrame, even if empty
if dwas_df is None or not isinstance(dwas_df, pd.DataFrame):
print("Warning: Invalid or missing DWA DataFrame. Proceeding without DWA data.")
dwas_df = pd.DataFrame(
columns=["onetsoc_code", "task_id", "dwa_title"]
) # Ensure it's an empty DF
print("Starting data processing...")
# --- 1. Handle Frequency (FT) ---
freq_df = ratings_df[ratings_df["scale_id"] == "FT"].copy()
if not freq_df.empty:
freq_pivot = freq_df.pivot_table(
index=["onetsoc_code", "task_id"],
columns="category",
values="data_value",
fill_value=0,
)
freq_pivot.columns = [
f"frequency_category_{int(col)}" for col in freq_pivot.columns
]
print(f"Processed Frequency data. Shape: {freq_pivot.shape}")
else:
print("No Frequency (FT) data found.")
# Create an empty DataFrame with the multi-index to allow merging later
idx = pd.MultiIndex(
levels=[[], []], codes=[[], []], names=["onetsoc_code", "task_id"]
)
freq_pivot = pd.DataFrame(index=idx)
# --- 2. Handle Importance (IM, IJ) ---
imp_df = ratings_df[ratings_df["scale_id"].isin(["IM", "IJ"])].copy()
if not imp_df.empty:
imp_avg = (
imp_df.groupby(["onetsoc_code", "task_id"])["data_value"]
.mean()
.reset_index()
)
imp_avg.rename(columns={"data_value": "importance_average"}, inplace=True)
print(f"Processed Importance data. Shape: {imp_avg.shape}")
else:
print("No Importance (IM, IJ) data found.")
imp_avg = pd.DataFrame(
columns=["onetsoc_code", "task_id", "importance_average"]
)
# --- 3. Handle Relevance (RT) ---
rel_df = ratings_df[ratings_df["scale_id"] == "RT"].copy()
if not rel_df.empty:
rel_avg = (
rel_df.groupby(["onetsoc_code", "task_id"])["data_value"]
.mean()
.reset_index()
)
rel_avg.rename(columns={"data_value": "relevance_average"}, inplace=True)
print(f"Processed Relevance data. Shape: {rel_avg.shape}")
else:
print("No Relevance (RT) data found.")
rel_avg = pd.DataFrame(columns=["onetsoc_code", "task_id", "relevance_average"])
# --- 4. Process DWAs ---
if dwas_df is not None and not dwas_df.empty and "dwa_title" in dwas_df.columns:
print("Processing DWA data...")
# Group DWAs by task_id and aggregate titles into a list
dwas_grouped = (
dwas_df.groupby(["onetsoc_code", "task_id"])["dwa_title"]
.apply(list)
.reset_index()
) #
dwas_grouped.rename(
columns={"dwa_title": "dwas"}, inplace=True
) # Rename column to 'dwas'
print(f"Processed DWA data. Shape: {dwas_grouped.shape}")
else:
print("No valid DWA data found or provided for processing.")
dwas_grouped = None # Set to None if no DWAs
# --- 5. Get Base Task/Occupation Info ---
base_cols = [
"onetsoc_code",
"task_id",
"task",
"occupation_title",
"occupation_description",
]
# Check if base columns exist in ratings_df
missing_base_cols = [col for col in base_cols if col not in ratings_df.columns]
if missing_base_cols:
print(
f"Error: Missing base info columns in ratings_df: {missing_base_cols}. Cannot proceed."
)
return None
if not ratings_df.empty:
base_info = (
ratings_df[base_cols]
.drop_duplicates()
.set_index(["onetsoc_code", "task_id"])
)
print(f"Extracted base info. Shape: {base_info.shape}")
else:
print("Cannot extract base info from empty ratings DataFrame.")
# Create an empty df with index to avoid errors later if possible
idx = pd.MultiIndex(
levels=[[], []], codes=[[], []], names=["onetsoc_code", "task_id"]
)
base_info = pd.DataFrame(
index=idx,
columns=[
col for col in base_cols if col not in ["onetsoc_code", "task_id"]
],
)
# --- 6. Merge Processed Data ---
print("Merging processed data...")
# Start with base_info, which should have the index ['onetsoc_code', 'task_id']
final_df = base_info.merge(
freq_pivot, left_index=True, right_index=True, how="left"
)
# Reset index before merging non-indexed dfs
final_df = final_df.reset_index()
# Merge averages - check if they are not empty before merging
if not imp_avg.empty:
final_df = final_df.merge(imp_avg, on=["onetsoc_code", "task_id"], how="left")
else:
final_df["importance_average"] = np.nan # Add column if imp_avg was empty
if not rel_avg.empty:
final_df = final_df.merge(rel_avg, on=["onetsoc_code", "task_id"], how="left")
else:
final_df["relevance_average"] = np.nan # Add column if rel_avg was empty
# Merge DWAs if available
if dwas_grouped is not None and not dwas_grouped.empty:
final_df = final_df.merge(
dwas_grouped, on=["onetsoc_code", "task_id"], how="left"
) # Merge the dwas list
# Fill NaN in 'dwas' column (for tasks with no DWAs) with empty lists
# Check if 'dwas' column exists before applying function
if "dwas" in final_df.columns:
final_df["dwas"] = final_df["dwas"].apply(
lambda x: x if isinstance(x, list) else []
) # Ensure tasks without DWAs get []
else:
print("Warning: 'dwas' column not created during merge.")
final_df["dwas"] = [
[] for _ in range(len(final_df))
] # Add empty list column
else:
# Add an empty 'dwas' column if no DWA data was processed or merged
final_df["dwas"] = [[] for _ in range(len(final_df))]
print(f"Final merged data shape: {final_df.shape}")
# Convert DataFrame to list of dictionaries for JSON output
# Handle potential NaN values during JSON conversion
# Replace numpy NaN with Python None for JSON compatibility
final_df = final_df.replace({np.nan: None})
result_list = final_df.to_dict(orient="records")
return result_list
# --- Output ---
def write_to_json(data, output_path):
"""
Writes the processed data to a JSON file.
Args:
data (list): The list of dictionaries to write.
output_path (str): Path to the output JSON file.
"""
if data is None:
print("No data to write to JSON.")
return
if not isinstance(data, list):
print(
f"Error: Data to write is not a list (type: {type(data)}). Cannot write to JSON."
)
return
# Create directory if it doesn't exist
output_dir = os.path.dirname(output_path)
if output_dir and not os.path.exists(output_dir):
try:
os.makedirs(output_dir)
print(f"Created output directory: {output_dir}")
except OSError as e:
print(f"Error creating output directory {output_dir}: {e}")
return # Exit if cannot create directory
try:
with open(output_path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=4, ensure_ascii=False)
print(f"Successfully wrote enriched data to {output_path}")
except IOError as e:
print(f"Error writing JSON file to {output_path}: {e}")
except TypeError as e:
print(f"Error during JSON serialization: {e}. Check data types.")
except Exception as e:
print(f"An unexpected error occurred during JSON writing: {e}")
# --- Main Execution ---
if __name__ == "__main__":
print("Starting O*NET Task Ratings & DWAs Enrichment Script...")
# 1. Fetch data
ratings_data_df, dwas_data_df = fetch_data_from_db(DB_FILE) # Fetch both datasets
# 2. Process data
# Proceed only if ratings_data_df is a valid DataFrame (even if empty)
# dwas_data_df can be None or empty, handled inside process function
if isinstance(ratings_data_df, pd.DataFrame):
enriched_data = process_task_ratings_with_dwas(
ratings_data_df, dwas_data_df
) # Pass both dataframes
# 3. Write output
if (
enriched_data is not None
): # Check if processing returned data (even an empty list is valid)
write_to_json(enriched_data, OUTPUT_FILE)
else:
print("Data processing failed or returned None. No output file generated.")
else:
print(
"Data fetching failed or returned invalid type for ratings data. Script terminated."
)
print("Script finished.")

81
pipeline/aggregate.py Normal file
View file

@ -0,0 +1,81 @@
from .utils import OCCUPATION_MAJOR_CODES
import pandas as pd
def create_task_summary_by_occupation_df(df_tasks: pd.DataFrame, oesm_df: pd.DataFrame) -> pd.DataFrame:
# --- OESM Wage Bill Calculation ---
df_oesm_with_bill = oesm_df.copy()
df_oesm_with_bill.rename(columns={'OCC_CODE': 'onetsoc_code'}, inplace=True)
# Convert key columns to numeric, handling potential errors
df_oesm_with_bill['TOT_EMP'] = pd.to_numeric(df_oesm_with_bill['TOT_EMP'], errors='coerce')
df_oesm_with_bill['A_MEAN'] = pd.to_numeric(df_oesm_with_bill['A_MEAN'], errors='coerce')
df_oesm_with_bill.dropna(subset=['TOT_EMP', 'A_MEAN', 'onetsoc_code'], inplace=True)
# Calculate the wage bill for each occupation
df_oesm_with_bill['wage_bill'] = df_oesm_with_bill['TOT_EMP'] * df_oesm_with_bill['A_MEAN']
oesm_lookup = df_oesm_with_bill.set_index('onetsoc_code')
summary_data = []
# Assuming df_tasks has an 'onetsoc_code' column with the full SOC code
unique_soc_codes = df_tasks['onetsoc_code'].unique()
for code in unique_soc_codes:
occ_df = df_tasks[df_tasks['onetsoc_code'] == code]
total_tasks_in_occ = len(occ_df)
not_remote_count = len(occ_df[occ_df['remote_status'] != 'remote'])
remote_df = occ_df[occ_df['remote_status'] == 'remote']
remote_estimable_count = len(remote_df[remote_df['estimable']])
remote_not_estimable_count = len(remote_df[~remote_df['estimable']])
try:
# O*NET codes (e.g., 11-1011.03) are more specific than OESM SOC codes (e.g., 11-1011).
# We strip the suffix from the O*NET code to find the corresponding wage data.
soc_code_for_lookup = code.split('.')[0]
wage_bill = oesm_lookup.loc[soc_code_for_lookup, 'wage_bill']
label = oesm_lookup.loc[soc_code_for_lookup, 'OCC_TITLE']
except KeyError:
wage_bill = 0
label = "Unknown"
summary_data.append({
'onetsoc_code': code,
'occupation_label': label,
'wage_bill': wage_bill,
'count_not_remote': not_remote_count,
'count_remote_estimable': remote_estimable_count,
'count_remote_not_estimable': remote_not_estimable_count,
'total_tasks': total_tasks_in_occ
})
return pd.DataFrame(summary_data)
def aggregate_task_summary_by_major_code(summary_df: pd.DataFrame) -> pd.DataFrame:
df_agg = summary_df.copy()
df_agg['onetsoc_major_code'] = df_agg['onetsoc_code'].str[:2]
aggregation = {
'wage_bill': 'sum',
'count_not_remote': 'sum',
'count_remote_estimable': 'sum',
'count_remote_not_estimable': 'sum',
'total_tasks': 'sum'
}
major_summary = df_agg.groupby('onetsoc_major_code').agg(aggregation).reset_index()
major_summary['occupation_label'] = major_summary['onetsoc_major_code'].map(OCCUPATION_MAJOR_CODES)
# Reorder columns to match original output format
major_summary = major_summary[[
'onetsoc_major_code',
'occupation_label',
'wage_bill',
'count_not_remote',
'count_remote_estimable',
'count_remote_not_estimable',
'total_tasks'
]]
return major_summary

225
pipeline/classification.py Normal file
View file

@ -0,0 +1,225 @@
from pathlib import Path
import pandas as pd
from .logger import logger
from .utils import enrich
import json
ALLOWED_UNITS = [
"minute",
"hour",
"day",
"week",
"month",
"trimester",
"semester",
"year",
]
ESTIMABLE_CLASSIFICATION_VERSION = "old_version"
TIME_ESTIMATES_GENERATION_VERSION = "old_version"
def classify_tasks_as_estimable(cache_dir: Path, df_to_process: pd.DataFrame, bust: bool = False) -> pd.DataFrame:
CACHE_PATH = cache_dir / f"task_estimability.{ESTIMABLE_CLASSIFICATION_VERSION}.parquet"
if CACHE_PATH.exists() and not bust:
logger.info(f"Loading cached task estimability from {CACHE_PATH}")
return pd.read_parquet(CACHE_PATH)
logger.info("Enriching tasks with estimability classification.")
df_unique_tasks = df_to_process.drop_duplicates(subset=['task']).copy()
logger.info(f"Found {len(df_unique_tasks)} unique remote tasks to classify.")
if df_unique_tasks.empty:
raise ValueError("No unique tasks to classify.")
results = enrich(
model="gpt-4.1-mini",
rpm=5000,
messages_to_process=[
[
{"role": "system", "content": """
Classify the provided O*NET task into one of these categories:
- ATOMIC (schedulable): A single, clearly-bounded activity, typically lasting minutes, hours, or a few days.
- ONGOING-CONSTRAINT (background role/ethical rule): A continuous responsibility or behavioural norm with no schedulable duration (e.g., follow confidentiality rules, serve as department head).
""".strip()},
{"role": "user", "content": f"Task: {row.task}"},
]
for row in df_unique_tasks.itertuples()
],
schema={
"name": "estimability_classification",
"schema": {
"type": "object",
"properties": {"task_category": {"type": "string", "enum": ["ATOMIC", "ONGOING-CONSTRAINT"]}},
"required": ["task_category"],
"additionalProperties": False
}
},
chunk_size=300,
)
if not results or len(results) != len(df_unique_tasks):
raise ValueError(f"Task estimability classification failed or returned mismatched number of results. Expected {len(df_unique_tasks)}, got {len(results) if results else 0}.")
classifications = []
for index, response in enumerate(results):
task_label = df_unique_tasks.iloc[index]['task']
task_category_flag = None
if response is None:
logger.warning(f"API call failed for task (enrich returned None): '{task_label}'")
else:
try:
content_str = response.choices[0].message.content
if not content_str:
raise ValueError("No content found in the response message")
data = json.loads(content_str)
if 'task_category' in data and isinstance(data['task_category'], str):
task_category_flag = data['task_category']
else:
logger.warning(f"Invalid or missing 'task_category' payload for task '{task_label}'. Data: '{data}'")
except (json.JSONDecodeError, AttributeError, KeyError, IndexError, ValueError) as e:
logger.warning(f"Could not parse response for task '{task_label}'. Error: {e}. Response: {response}")
classifications.append({
'task': task_label,
'estimable': task_category_flag == 'ATOMIC'
})
classification_df = pd.DataFrame(classifications)
logger.info(f"Finished classification. Got {classification_df['estimable'].notna().sum()} successful classifications out of {len(df_unique_tasks)} unique tasks.")
logger.info(f"Saving task estimability classifications to {CACHE_PATH}")
classification_df.to_parquet(CACHE_PATH)
return classification_df
def generate_time_estimates_for_tasks(cache_dir: Path, df_to_process: pd.DataFrame, bust: bool = False) -> pd.DataFrame:
CACHE_PATH = cache_dir / f"task_estimates.{TIME_ESTIMATES_GENERATION_VERSION}.parquet"
if CACHE_PATH.exists() and not bust:
logger.info(f"Loading cached task estimates from {CACHE_PATH}")
return pd.read_parquet(CACHE_PATH)
logger.info("Enriching tasks with time estimates.")
if df_to_process.empty:
raise ValueError("No tasks to process for estimates.")
results = enrich(
model="gpt-4.1-mini",
rpm=5000,
messages_to_process=[
[
{
"role": "system",
"content": """
You are an expert assistant evaluating the time required for job tasks. Your goal is to estimate the 'effective time' range needed for a skilled human to complete the following job task **remotely**, without supervision
'Effective time' is the active, focused work duration required to complete the task. Crucially, **exclude all waiting periods, delays, or time spent on other unrelated activities**. Think of it as the continuous, productive time investment needed if the worker could pause and resume instantly without cost.
Provide a lower and upper bound estimate for the 'effective time'. These bounds should capture the time within which approximately 80% of instances of performing this specific task are typically completed by a qualified individual.
Base your estimate on the provided task and the associated occupation and occupation description. Your estimate must be in one the allowed units: minute, hour, day, week, month, trimester, semester, year.""".strip()
},
{
"role": "user",
"content": f"{row.task} done by {row.occupation_title} ({row.occupation_description})"
}
]
for row in df_to_process.itertuples()
],
schema= {
"name": "estimate_time",
"strict": True,
"schema": {
"type": "object",
"properties": {
"lower_bound_estimate": {
"type": "object",
"properties": {
"quantity": {
"type": "number",
"description": "The numerical value for the lower bound of the estimate.",
},
"unit": {
"type": "string",
"enum": ALLOWED_UNITS,
"description": "The unit of time for the lower bound.",
},
},
"required": ["quantity", "unit"],
"additionalProperties": False,
},
"upper_bound_estimate": {
"type": "object",
"properties": {
"quantity": {
"type": "number",
"description": "The numerical value for the upper bound of the estimate.",
},
"unit": {
"type": "string",
"enum": ALLOWED_UNITS,
"description": "The unit of time for the upper bound.",
},
},
"required": ["quantity", "unit"],
"additionalProperties": False,
},
},
"required": ["lower_bound_estimate", "upper_bound_estimate"],
"additionalProperties": False,
},
},
chunk_size=200,
)
if not results or len(results) != len(df_to_process):
raise ValueError(f"API call for task estimates failed or returned mismatched number of results. "
f"Expected {len(df_to_process)}, got {len(results) if results else 0}.")
estimates = []
for index, response in enumerate(results):
row = df_to_process.iloc[index]
task_info = f"O*NET: {row.onetsoc_code}, Task ID: {row.task_id}"
lb_qty, lb_unit, ub_qty, ub_unit = None, None, None, None
if response is None:
logger.warning(f"API call failed for task (enrich returned None): {task_info}")
else:
try:
content_str = response.choices[0].message.content
if not content_str:
raise ValueError("No content found in the response message")
data = json.loads(content_str)
lb_qty = data['lower_bound_estimate']['quantity']
lb_unit = data['lower_bound_estimate']['unit']
ub_qty = data['upper_bound_estimate']['quantity']
ub_unit = data['upper_bound_estimate']['unit']
except Exception as e:
logger.warning(f"Could not parse valid estimate for task {task_info}. Error: {e}. Response: {response}")
lb_qty, lb_unit, ub_qty, ub_unit = None, None, None, None # Reset on failure
estimates.append({
'onetsoc_code': row.onetsoc_code,
'task_id': row.task_id,
'lb_estimate_qty': lb_qty,
'lb_estimate_unit': lb_unit,
'ub_estimate_qty': ub_qty,
'ub_estimate_unit': ub_unit
})
estimates_df = pd.DataFrame(estimates)
logger.info(f"Finished estimates. Got {estimates_df['lb_estimate_qty'].notna().sum()} successful estimates out of {len(df_to_process)} tasks.")
logger.info(f"Saving task estimates to {CACHE_PATH}")
estimates_df.to_parquet(CACHE_PATH)
return estimates_df

132
pipeline/fetchers.py Normal file
View file

@ -0,0 +1,132 @@
import sqlite3
import pandas as pd
import requests
import io
import zipfile
import yaml
from pathlib import Path
from .logger import logger
from typing import Tuple, Dict
ONET_VERSION = "29_1"
ONET_URL = f"https://www.onetcenter.org/dl_files/database/db_{ONET_VERSION}_mysql.zip"
def fetch_onet_database(cache_dir: Path) -> sqlite3.Connection:
DB_PATH = cache_dir / f"onet_{ONET_VERSION}.db"
if DB_PATH.exists():
logger.info(f"Using cached O*NET database: {DB_PATH}")
return sqlite3.connect(DB_PATH)
logger.info(f"Downloading O*NET database from {ONET_URL}")
response = requests.get(ONET_URL, stream=True, headers={
"User-Agent": "econ-agent/1.0"
})
response.raise_for_status()
conn = sqlite3.connect(DB_PATH)
conn.executescript("""
PRAGMA journal_mode = OFF;
PRAGMA synchronous = 0;
PRAGMA cache_size = 1000000;
PRAGMA locking_mode = EXCLUSIVE;
PRAGMA temp_store = MEMORY;
PRAGMA foreign_keys = ON;
""")
zip_content = response.content
with zipfile.ZipFile(io.BytesIO(zip_content)) as z:
sql_scripts = []
for filename in sorted(z.namelist()):
if filename.endswith(".sql"):
sql_scripts.append(z.read(filename).decode('utf-8'))
if not sql_scripts:
raise RuntimeError("No SQL files found in the O*NET zip archive.")
logger.info("Executing SQL files in alphabetical order (single transaction mode)")
full_script = "BEGIN TRANSACTION;\n" + "\n".join(sql_scripts) + "\nCOMMIT;"
conn.executescript(full_script)
conn.executescript("""
PRAGMA journal_mode = WAL;
PRAGMA synchronous = NORMAL;
PRAGMA locking_mode = NORMAL;
PRAGMA temp_store = DEFAULT;
PRAGMA foreign_keys = ON;
PRAGMA optimize;
""")
conn.execute("VACUUM;")
conn.commit()
return conn
def fetch_oesm_data(cache_dir: Path) -> pd.DataFrame:
VERSION = "23"
URL = f"https://www.bls.gov/oes/special-requests/oesm{VERSION}nat.zip"
DATA_PATH = cache_dir / "oesm.parquet"
if DATA_PATH.exists():
logger.info(f"Using cached OESM data: {DATA_PATH}")
return pd.read_parquet(DATA_PATH)
logger.info(f"Downloading OESM data from {URL}")
headers = {'User-Agent': 'econ-agent/1.0'}
response = requests.get(URL, headers=headers)
response.raise_for_status()
zip_content = response.content
logger.info(f"Creating new OESM data cache: {DATA_PATH}")
with zipfile.ZipFile(io.BytesIO(zip_content)) as z:
with z.open(f"oesm{VERSION}national.xlsx") as f:
df = pd.read_excel(f, engine='openpyxl', na_values=['*', '#'])
df.to_parquet(DATA_PATH)
logger.info(f"Saved OESM data to cache: {DATA_PATH}")
return df
def fetch_epoch_remote_data(cache_dir: Path) -> pd.DataFrame:
URL = "https://drive.google.com/uc?export=download&id=1GrHhuYIgaCCgo99dZ_40BWraz-fzo76r"
DATA_PATH = cache_dir / f"epoch_remote_latest.parquet"
if DATA_PATH.exists():
logger.info(f"Using cached EPOCH remote data: {DATA_PATH}")
return pd.read_parquet(DATA_PATH)
logger.info(f"Downloading EPOCH remote data from Google Drive: {URL}")
session = requests.Session()
session.headers.update({"User-Agent": "econ-agent/1.0"})
response = session.get(URL, stream=True)
response.raise_for_status()
csv_content = response.content
logger.info(f"Creating new EPOCH remote data cache: {DATA_PATH}")
df = pd.read_csv(io.BytesIO(csv_content))
df.to_parquet(DATA_PATH)
return df
def fetch_metr_data(cache_dir: Path) -> Dict:
URL = "https://metr.org/assets/benchmark_results.yaml"
DATA_PATH = cache_dir / "metr_benchmark_results.yaml"
if DATA_PATH.exists():
logger.info(f"Using cached METR data: {DATA_PATH}")
with open(DATA_PATH, "r") as f:
return yaml.safe_load(f)
logger.info(f"Downloading METR data from {URL}")
headers = {"User-Agent": "econ-agent/1.0"}
response = requests.get(URL, headers=headers)
response.raise_for_status()
yaml_content = response.content
logger.info(f"Creating new METR data cache: {DATA_PATH}")
with open(DATA_PATH, "wb") as f:
f.write(yaml_content)
return yaml.safe_load(yaml_content)

View file

@ -0,0 +1,15 @@
from .estimate_histplot import generate_estimate_histplot
from .estimates_spread_per_occupation import generate_estimate_spread_per_occupation
from .estimates_lower_vs_upper_scatter import generate_estimates_lower_vs_upper_scatter
from .sequential_coherence_cdf import plot_sequential_coherence_cdf
from .projected_automatable_wage_bill import generate_projected_automatable_wage_bill
from .projected_task_automation import generate_projected_task_automation_plot
GENERATORS = [
generate_estimate_histplot,
generate_estimate_spread_per_occupation,
generate_estimates_lower_vs_upper_scatter,
#plot_sequential_coherence_cdf,
generate_projected_automatable_wage_bill,
generate_projected_task_automation_plot,
]

View file

@ -0,0 +1,32 @@
from pathlib import Path
from typing import Generator
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from ..utils import style_plot
def generate_estimate_histplot(output_dir: Path, df: pd.DataFrame, **kwargs) -> Generator[Path]:
"""
Generates a styled histogram of the distribution of midpoint time estimates.
"""
style_plot()
OUTPUT_PATH = output_dir / "estimate_distribution_histplot.png"
fig, ax = plt.subplots()
sns.histplot(
data=df,
x='estimate_midpoint',
log_scale=True,
ax=ax
)
ax.set_xlabel("Task Time (minutes, log scale)")
ax.set_ylabel("Number of Tasks")
ax.set_title("Distribution of Time Estimates for Atomic Tasks")
plt.tight_layout()
plt.savefig(OUTPUT_PATH)
plt.close(fig)
yield OUTPUT_PATH

View file

@ -0,0 +1,56 @@
from pathlib import Path
from typing import Generator
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from ..utils import OCCUPATION_MAJOR_CODES, style_plot
def generate_estimates_lower_vs_upper_scatter(output_dir: Path, df: pd.DataFrame, **kwargs) -> Generator[Path]:
"""
Generates a styled scatter plot of lower-bound vs upper-bound time estimates for tasks.
"""
style_plot()
OUTPUT_PATH = output_dir / "estimates_lower_vs_upper_scatter.png"
plot_df = df.copy()
# Replace onetsoc_major codes with their corresponding labels for the plot legend
plot_df['onetsoc_major'] = plot_df['onetsoc_major'].map(OCCUPATION_MAJOR_CODES)
fig, ax = plt.subplots(figsize=(12, 10))
sns.scatterplot(
data=plot_df,
x='lb_estimate_in_minutes',
y='ub_estimate_in_minutes',
alpha=0.3,
edgecolor=None,
hue="onetsoc_major",
ax=ax
)
# 45° reference line (y=x)
lims = (
min(df['lb_estimate_in_minutes'].min(), df['ub_estimate_in_minutes'].min()),
max(df['lb_estimate_in_minutes'].max(), df['ub_estimate_in_minutes'].max())
)
lims = (lims[0] * 0.9, lims[1] * 1.1)
ax.plot(lims, lims, color='black', linestyle='--', linewidth=1, zorder=0)
# Optional helper lines for ratios
for k in [2, 10, 100]:
ax.plot(lims, [k*l for l in lims],
linestyle=':', color='grey', linewidth=1, zorder=0)
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel('Lower-bound (min, log scale)')
ax.set_ylabel('Upper-bound (min, log scale)')
ax.set_title('Lower vs Upper Estimates for All Tasks')
ax.legend(title="Occupation Major Group", bbox_to_anchor=(1.02, 1), loc='upper left')
plt.tight_layout()
plt.savefig(OUTPUT_PATH, bbox_inches='tight')
plt.close(fig)
yield OUTPUT_PATH

View file

@ -0,0 +1,39 @@
from pathlib import Path
from typing import Generator
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from ..utils import OCCUPATION_MAJOR_CODES, style_plot
def generate_estimate_spread_per_occupation(output_dir: Path, df: pd.DataFrame, **kwargs) -> Generator[Path]:
"""
Generates a styled boxplot of the estimate range spread per major occupation group.
"""
style_plot()
OUTPUT_PATH = output_dir / "estimates_spread_per_occupation.png"
fig, ax = plt.subplots(figsize=(10, 12))
sns.boxplot(
data=df,
x='onetsoc_major',
y='estimate_range',
showfliers=False,
ax=ax
)
ax.set_yscale('log')
ax.set_xlabel('Occupation')
ax.set_ylabel('Range (upper-lower, minutes)')
ax.set_title('Spread of time-range estimates per occupation')
# Get occupation labels from codes for x-axis ticks
labels = [OCCUPATION_MAJOR_CODES.get(code.get_text(), code.get_text()) for code in ax.get_xticklabels()]
ax.set_xticklabels(labels, rotation=60, ha='right')
plt.tight_layout()
plt.savefig(OUTPUT_PATH)
plt.close(fig)
yield OUTPUT_PATH

View file

@ -0,0 +1,229 @@
from pathlib import Path
from typing import Generator, Dict, Tuple, Optional
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from scipy.stats import linregress
from datetime import datetime
from ..utils import style_plot, LIME
def _generate_wage_projection_data(
metr_results: Dict,
df_with_wages: pd.DataFrame,
percentile_key: str,
doubling_time_modifier: float,
) -> Optional[Tuple[pd.DataFrame, pd.DataFrame, float]]:
"""
Generates wage projection data for different AI progress scenarios.
Args:
metr_results: The METR benchmark data.
df_with_wages: DataFrame containing tasks with their estimated wage value.
percentile_key: The percentile to use from METR data (e.g., 'p50_horizon_length').
doubling_time_modifier: Multiplier for the doubling time (e.g., 1.0 for baseline,
0.5 for optimistic, 2.0 for pessimistic).
Returns:
A tuple of (metr_df, projection_df, doubling_time_days), or None if data is insufficient.
"""
all_model_data = []
for model_name, data in metr_results.get("results", {}).items():
for agent_name, agent_data in data.get("agents", {}).items():
release_date_str = data.get("release_date")
horizon = agent_data.get(percentile_key, {}).get("estimate")
if release_date_str and horizon is not None:
all_model_data.append({
"release_date": release_date_str,
"horizon_minutes": horizon,
})
if not all_model_data:
return None
metr_df = pd.DataFrame(all_model_data).sort_values("release_date").reset_index(drop=True)
metr_df['release_date'] = pd.to_datetime(metr_df['release_date'])
metr_df = metr_df[metr_df['horizon_minutes'] > 0].copy()
if len(metr_df) < 2:
return None
metr_df['days_since_start'] = (metr_df['release_date'] - metr_df['release_date'].min()).dt.days
log_y = np.log(metr_df['horizon_minutes'])
slope, intercept, r_value, _, _ = linregress(metr_df['days_since_start'], log_y)
# Apply the scenario modifier to the doubling time
base_doubling_time_days = np.log(2) / slope
modified_doubling_time_days = base_doubling_time_days * doubling_time_modifier
modified_slope = np.log(2) / modified_doubling_time_days
start_date = metr_df['release_date'].min()
future_dates = pd.to_datetime(pd.date_range(start=start_date, end="2035-01-01", freq="ME"))
future_days = (future_dates - start_date).days.to_numpy()
projected_log_horizon = intercept + modified_slope * future_days
projected_horizon_minutes = np.exp(projected_log_horizon)
projection_df = pd.DataFrame({
"date": future_dates,
"projected_coherence_minutes": projected_horizon_minutes,
})
# Calculate the total wage bill of tasks automated over time
for bound in ["lb", "mid", "ub"]:
col_name = 'estimate_midpoint' if bound == 'mid' else f'{bound}_estimate_in_minutes'
projection_df[f"automatable_wage_bill_{bound}"] = projection_df["projected_coherence_minutes"].apply(
lambda h: df_with_wages.loc[df_with_wages[col_name] <= h, 'wage_per_task'].sum()
)
# Also calculate for the actual METR data points for plotting
metr_df["automatable_wage_bill_mid"] = metr_df["horizon_minutes"].apply(
lambda h: df_with_wages.loc[df_with_wages['estimate_midpoint'] <= h, 'wage_per_task'].sum()
)
return metr_df, projection_df, modified_doubling_time_days
def _plot_scenario(ax, projection_df, metr_df, label, color, line_style='-'):
"""Helper function to draw a single projection scenario on a given axis."""
# Plot the projected wage bill
ax.plot(
projection_df["date"],
projection_df["automatable_wage_bill_mid"],
label=label,
color=color,
linewidth=2.5,
linestyle=line_style,
zorder=3
)
# Plot the shaded range for lower/upper bounds
ax.fill_between(
projection_df["date"],
projection_df["automatable_wage_bill_lb"],
projection_df["automatable_wage_bill_ub"],
color=color,
alpha=0.15,
zorder=2
)
# Plot the actual METR data points against the wage bill
ax.scatter(
metr_df['release_date'],
metr_df['automatable_wage_bill_mid'],
color=color,
edgecolor='black',
s=60,
zorder=4,
label=f"Model Capabilities (P50)"
)
def generate_projected_automatable_wage_bill(
output_dir: Path,
df: pd.DataFrame,
task_summary_by_occupation_df: pd.DataFrame,
metr_results: Dict,
**kwargs,
) -> Generator[Path, None, None]:
"""
Generates a plot projecting the automatable wage bill under different
AI progress scenarios (optimistic, baseline, pessimistic).
"""
style_plot()
OUTPUT_PATH = output_dir / "projected_automatable_wage_bill_sensitivity.png"
# 1. Calculate wage_per_task for each occupation
wage_bill_info = task_summary_by_occupation_df[['onetsoc_code', 'wage_bill', 'total_tasks']].copy()
wage_bill_info['wage_per_task'] = wage_bill_info['wage_bill'] / wage_bill_info['total_tasks']
wage_bill_info.replace([np.inf, -np.inf], 0, inplace=True) # Avoid division by zero issues
wage_bill_info.drop(columns=['wage_bill', 'total_tasks'], inplace=True)
# 2. Merge wage_per_task into the main task dataframe
df_with_wages = pd.merge(df, wage_bill_info, on='onetsoc_code', how='left')
df_with_wages['wage_per_task'].fillna(0, inplace=True)
# 3. Generate data for all three scenarios
scenarios = {
"Optimistic": {"modifier": 0.5, "color": "tab:green", "style": "--"},
"Baseline": {"modifier": 1.0, "color": LIME['600'], "style": "-"},
"Pessimistic": {"modifier": 2.0, "color": "tab:red", "style": ":"},
}
projection_results = {}
for name, config in scenarios.items():
result = _generate_wage_projection_data(metr_results, df_with_wages, 'p50_horizon_length', config['modifier'])
if result:
projection_results[name] = result
if not projection_results:
print("Warning: Could not generate any projection data. Skipping wage bill plot.")
return
# 4. Create the plot
fig, ax = plt.subplots(figsize=(14, 9))
# We only need to plot the scatter points once, let's use the baseline ones.
if "Baseline" in projection_results:
metr_df, _, _ = projection_results["Baseline"]
ax.scatter(
metr_df['release_date'],
metr_df['automatable_wage_bill_mid'],
color='black',
s=80,
zorder=5,
label=f"Model Capabilities (P50)"
)
legend_lines = []
for name, (metr_df, proj_df, doubling_time) in projection_results.items():
config = scenarios[name]
ax.plot(
proj_df["date"],
proj_df["automatable_wage_bill_mid"],
color=config['color'],
linestyle=config['style'],
linewidth=2.5,
zorder=3
)
ax.fill_between(
proj_df["date"],
proj_df["automatable_wage_bill_lb"],
proj_df["automatable_wage_bill_ub"],
color=config['color'],
alpha=0.15,
zorder=2
)
# Create a custom line for the legend
line = plt.Line2D([0], [0], color=config['color'], linestyle=config['style'], lw=2.5,
label=f'{name} (Doubling Time: {doubling_time:.0f} days)')
legend_lines.append(line)
# 5. Styling and annotations
ax.set_title("Projected Automatable Wage Bill (P50 Coherence)", fontsize=18, pad=20)
ax.set_xlabel("Year", fontsize=12)
ax.set_ylabel("Automatable Annual Wage Bill (Trillions of USD)", fontsize=12)
# Format Y-axis to show trillions
def trillions_formatter(x, pos):
return f'${x / 1e12:.1f}T'
ax.yaxis.set_major_formatter(mticker.FuncFormatter(trillions_formatter))
total_wage_bill = df_with_wages['wage_per_task'].sum()
ax.set_ylim(0, total_wage_bill * 1.05)
if "Baseline" in projection_results:
_, proj_df, _ = projection_results["Baseline"]
ax.set_xlim(datetime(2022, 1, 1), proj_df["date"].max())
# Create the legend from the custom lines and the scatter plot
scatter_legend = ax.get_legend_handles_labels()[0]
ax.legend(handles=legend_lines + scatter_legend, loc="upper left", fontsize=11)
ax.grid(True, which="both", linestyle="--", linewidth=0.5)
plt.tight_layout()
plt.savefig(OUTPUT_PATH)
plt.close(fig)
print(f"Generated sensitivity analysis plot: {OUTPUT_PATH}")
yield OUTPUT_PATH

View file

@ -0,0 +1,168 @@
from pathlib import Path
from typing import Generator, Dict, Tuple
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import linregress
from datetime import datetime
from ..utils import style_plot, LIME
def _generate_projection_data(
metr_results: Dict,
df: pd.DataFrame,
percentile_key: str,
) -> Tuple[pd.DataFrame, pd.DataFrame] | None:
"""
Generates projection data for a given percentile key (e.g., 'p50_horizon_length').
Returns a tuple of (metr_df_with_pct, projection_df), or None if data is insufficient.
"""
# 1. Process METR data to get all model performance over time for the given percentile
all_model_data = []
for model_name, data in metr_results.get("results", {}).items():
for agent_name, agent_data in data.get("agents", {}).items():
release_date_str = data.get("release_date")
horizon = agent_data.get(percentile_key, {}).get("estimate")
if release_date_str and horizon is not None:
unique_model_name = f"{model_name}-{agent_name}"
all_model_data.append({
"model": unique_model_name,
"release_date": release_date_str,
"horizon_minutes": horizon,
})
if not all_model_data:
print(f"Warning: No models with {percentile_key} found in METR data. Skipping.")
return None
metr_df = pd.DataFrame(all_model_data).sort_values("release_date").reset_index(drop=True)
metr_df['release_date'] = pd.to_datetime(metr_df['release_date'])
# 2. Perform log-linear regression on coherence over time
metr_df = metr_df[metr_df['horizon_minutes'] > 0].copy()
if len(metr_df) < 2:
print(f"Warning: Not enough data points for regression for {percentile_key}. Skipping.")
return None
metr_df['days_since_start'] = (metr_df['release_date'] - metr_df['release_date'].min()).dt.days
log_y = np.log(metr_df['horizon_minutes'])
x = metr_df['days_since_start']
slope, intercept, r_value, _, _ = linregress(x, log_y)
doubling_time_days = np.log(2) / slope
print(f"METR all models {percentile_key} trend: R^2 = {r_value**2:.2f}, Doubling time = {doubling_time_days:.1f} days")
# 3. Project coherence into the future
start_date = metr_df['release_date'].min()
future_dates = pd.to_datetime(pd.date_range(start=start_date, end="2035-01-01", freq="ME"))
future_days = (future_dates - start_date).days.to_numpy()
projected_log_horizon = intercept + slope * future_days
projected_horizon_minutes = np.exp(projected_log_horizon)
projection_df = pd.DataFrame({
"date": future_dates,
"projected_coherence_minutes": projected_horizon_minutes,
})
# 4. Calculate the percentage of tasks automated over time based on our estimates
total_tasks = len(df)
if total_tasks == 0:
return None
for bound in ["lb", "mid", "ub"]:
col_name = 'estimate_midpoint' if bound == 'mid' else f'{bound}_estimate_in_minutes'
projection_df[f"pct_automatable_{bound}"] = projection_df["projected_coherence_minutes"].apply(
lambda h: (df[col_name] <= h).sum() / total_tasks * 100
)
metr_df["pct_automatable_mid"] = metr_df["horizon_minutes"].apply(
lambda h: (df['estimate_midpoint'] <= h).sum() / total_tasks * 100
)
return metr_df, projection_df
def _plot_projection(ax, projection_df, metr_df, label, color, line_style='-'):
"""Helper function to draw a single projection on a given axis."""
# Plot the projected automation percentage
ax.plot(
projection_df["date"],
projection_df["pct_automatable_mid"],
label=f"Mid-point",
color=color,
linewidth=2.5,
linestyle=line_style,
zorder=3
)
ax.fill_between(
projection_df["date"],
projection_df["pct_automatable_lb"],
projection_df["pct_automatable_ub"],
color=color,
alpha=0.15,
label=f"Lower/upper bound range",
zorder=2
)
# Plot the actual METR data points
ax.scatter(
metr_df['release_date'],
metr_df['pct_automatable_mid'],
color=color,
edgecolor='black',
s=60,
zorder=4,
label=f"Model with {label[1:]}% success rate"
)
def generate_projected_task_automation_plot(
output_dir: Path,
metr_results: Dict,
df: pd.DataFrame,
**kwargs,
) -> Generator[Path, None, None]:
"""
Generates plots projecting task automation based on METR's p50 and p80
coherence data.
"""
style_plot()
p50_data = _generate_projection_data(metr_results, df, 'p50_horizon_length')
p80_data = _generate_projection_data(metr_results, df, 'p80_horizon_length')
# Plot P50 alone
if p50_data:
p50_metr_df, p50_proj_df = p50_data
fig, ax = plt.subplots(figsize=(12, 8))
_plot_projection(ax, p50_proj_df, p50_metr_df, "P50", LIME['600'])
ax.set_title("How long before sequential coherence stops being a bottleneck?", fontsize=16, pad=20)
ax.set_xlabel("Year")
ax.set_ylabel("% of task automatable (50% success rate)")
ax.set_ylim(0, 100.5)
ax.set_xlim(datetime(2022, 1, 1), p50_proj_df["date"].max())
ax.grid(True, which="both", linestyle="--", linewidth=0.5)
ax.legend(loc="upper left")
plt.tight_layout()
output_path = output_dir / "projected_task_automation_p50.png"
plt.savefig(output_path)
plt.close(fig)
yield output_path
# Plot P80 alone
if p80_data:
p80_metr_df, p80_proj_df = p80_data
fig, ax = plt.subplots(figsize=(12, 8))
_plot_projection(ax, p80_proj_df, p80_metr_df, "P80", 'tab:cyan')
ax.set_title("Projected Task Automation (P80 AI Coherence)", fontsize=16, pad=20)
ax.set_xlabel("Year")
ax.set_ylabel("% of Estimable Economic Tasks Automatable")
ax.set_ylim(0, 100.5)
ax.set_xlim(datetime(2022, 1, 1), p80_proj_df["date"].max())
ax.grid(True, which="both", linestyle="--", linewidth=0.5)
ax.legend(loc="upper left")
plt.tight_layout()
output_path = output_dir / "projected_task_automation_p80.png"
plt.savefig(output_path)
plt.close(fig)
yield output_path

View file

@ -0,0 +1,54 @@
from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
from ..utils import LIME, style_plot
def plot_sequential_coherence_cdf(output_dir: Path, df: pd.DataFrame, **kwargs):
style_plot()
output_path = output_dir / "sequential_coherence_cdf.png"
def cdf(series):
"""Helper function to calculate CDF data."""
s = series.sort_values().reset_index(drop=True)
# Calculate cumulative percentage
return s.values, ((s.index + 1) / len(s)) * 100
# Calculate CDF for lower, upper, and midpoint estimates
x_lb, y_lb = cdf(df['lb_estimate_in_minutes'])
x_ub, y_ub = cdf(df['ub_estimate_in_minutes'])
x_mid, y_mid = cdf(df['estimate_midpoint'])
# Create the plot
fig, ax = plt.subplots(figsize=(12, 7))
# Plot the CDFs as step plots
ax.step(x_lb, y_lb, where='post', color=LIME['300'], linewidth=1.8, linestyle='--', zorder=2, label='Lower bound estimate')
ax.step(x_ub, y_ub, where='post', color=LIME['900'], linewidth=1.8, linestyle=':', zorder=3, label='Upper bound estimate')
ax.step(x_mid, y_mid, where='post', color=LIME['600'], linewidth=2.2, zorder=4, label='Mid-point')
# --- Styling and Annotations ---
ax.set_xscale('log')
ax.set_ylim(0, 100)
ax.yaxis.set_major_formatter(mtick.PercentFormatter(decimals=0))
# Set titles and labels using the standard axes methods
ax.set_title("% of Tasks With Sequential Coherence ≤ X")
ax.set_xlabel("Sequential Coherence (X)")
ax.set_ylabel("Cumulative Percentage of Tasks")
# Define custom x-axis ticks and labels for better readability
ticks = [1, 5, 10, 30, 60, 120, 240, 480, 1440, 2880, 10080, 43200, 129600, 259200, 525600]
ticklabels = ['1 min', '5 min', '10 min', '30 min', '1 hr', '2 hr', '4 hr', '8 hr', '1 day', '2 days',
'1 wk', '30 days', '90 days', '180 days', '1 yr']
ax.set_xticks(ticks)
ax.set_xticklabels(ticklabels, rotation=45, ha='right')
ax.legend(loc='lower right')
# --- Save and close ---
plt.tight_layout()
plt.savefig(output_path, bbox_inches='tight')
plt.close(fig)
yield output_path

24
pipeline/logger.py Normal file
View file

@ -0,0 +1,24 @@
import logging
from logging.handlers import RotatingFileHandler
from rich.logging import RichHandler
LOGGER_NAME = "pipeline"
def setup_logging() -> logging.Logger:
# Set up Rich console handler
rich_handler = RichHandler(
level=logging.DEBUG,
show_time=True,
enable_link_path=True,
rich_tracebacks=True,
# omit_repeated_times=False,
)
logger = logging.getLogger(LOGGER_NAME)
logger.setLevel(logging.DEBUG)
logger.addHandler(rich_handler)
return logger
logger = setup_logging()

215
pipeline/runner.py Normal file
View file

@ -0,0 +1,215 @@
import sqlite3
import os
from .logger import logger
import pandas as pd
from dotenv import load_dotenv
from .fetchers import fetch_onet_database, fetch_oesm_data, fetch_epoch_remote_data, ONET_VERSION, fetch_metr_data
from .classification import classify_tasks_as_estimable, generate_time_estimates_for_tasks
from .generators import GENERATORS
from .aggregate import create_task_summary_by_occupation_df, aggregate_task_summary_by_major_code
from .utils import convert_to_minutes
import argparse
import platformdirs
import numpy as np
from pathlib import Path
class Runner:
onet_conn: sqlite3.Connection
oesm_df: pd.DataFrame
epoch_df: pd.DataFrame
metr_results: dict
def __init__(self, output_dir: Path | str, debug: bool, bust_estimability: bool, bust_estimates: bool):
if isinstance(output_dir, str):
output_dir = Path(output_dir).resolve()
output_dir.mkdir(parents=True, exist_ok=True)
self.output_dir = output_dir
self.intermediate_dir = self.output_dir / "intermediate"
self.intermediate_dir.mkdir(parents=True, exist_ok=True)
self.cache_dir = platformdirs.user_cache_path("econtai")
self.debug = debug
self.bust_estimability = bust_estimability
self.bust_estimates = bust_estimates
if debug:
os.environ["LITELLM_LOG"] = os.environ.get("LITELLM_LOG", "INFO")
def run(self):
load_dotenv()
self.onet_conn = fetch_onet_database(self.cache_dir)
self.oesm_df = fetch_oesm_data(self.cache_dir)
self.epoch_df = fetch_epoch_remote_data(self.cache_dir)
self.metr_results = fetch_metr_data(self.cache_dir)
self.df_tasks = self._create_df_tasks()
self.df_tasks['onetsoc_major'] = self.df_tasks['onetsoc_code'].str[:2]
df_to_process = self.df_tasks[
(self.df_tasks['importance_average'] > 3) &
(self.df_tasks['remote_status'] == 'remote')
].copy()
if self.debug:
df_to_process = df_to_process.head(10)
task_estimability_df = classify_tasks_as_estimable(self.cache_dir, df_to_process, bust=self.bust_estimability)
self.df_tasks = pd.merge(self.df_tasks, task_estimability_df, on='task', how='left')
self.df_tasks['estimable'] = self.df_tasks['estimable'].fillna(False)
self.df_tasks.to_parquet(self.intermediate_dir / "df_tasks.parquet")
df_to_process = pd.merge(df_to_process, task_estimability_df, on='task', how='left')
df_to_process['estimable'] = self.df_tasks['estimable'].fillna(False)
df_to_process = df_to_process[df_to_process['estimable']].copy()
task_estimates_df = generate_time_estimates_for_tasks(self.cache_dir, df_to_process, bust=self.bust_estimates)
df = pd.merge(df_to_process, task_estimates_df, on=['onetsoc_code', 'task_id'], how='left')
df['lb_estimate_in_minutes'] = df.apply(lambda row: convert_to_minutes(row['lb_estimate_qty'], row['lb_estimate_unit']), axis=1)
df['ub_estimate_in_minutes'] = df.apply(lambda row: convert_to_minutes(row['ub_estimate_qty'], row['ub_estimate_unit']), axis=1)
df['estimate_range'] = df.ub_estimate_in_minutes - df.lb_estimate_in_minutes
df['estimate_ratio'] = np.divide(df.ub_estimate_in_minutes, df.lb_estimate_in_minutes).replace([np.inf, -np.inf], None)
df['estimate_midpoint'] = (df.lb_estimate_in_minutes + df.ub_estimate_in_minutes) / 2
df.to_parquet(self.intermediate_dir / "estimable_tasks_with_estimates.parquet")
self.task_summary_by_occupation_df = create_task_summary_by_occupation_df(self.df_tasks, self.oesm_df)
self.task_summary_by_occupation_df.to_parquet(self.intermediate_dir / "task_summary_by_occupation.parquet")
self.task_summary_by_major_occupation_df = aggregate_task_summary_by_major_code(self.task_summary_by_occupation_df)
self.task_summary_by_major_occupation_df.to_parquet(self.intermediate_dir / "task_summary_by_major_occupation.parquet")
self._check_for_insanity(df)
for gen in GENERATORS:
for asset in gen(**{
"output_dir": self.output_dir,
"runner": self,
"df": df,
"task_summary_by_occupation_df": self.task_summary_by_occupation_df,
"task_summary_by_major_occupation_df": self.task_summary_by_major_occupation_df,
"df_tasks": self.df_tasks,
"oesm_df": self.oesm_df,
"metr_results": self.metr_results,
}):
logger.info(f"New asset: {asset}")
def _create_df_tasks(self) -> pd.DataFrame:
DATA_PATH = self.cache_dir / f"onet_{ONET_VERSION}_tasks_with_remote_status.parquet"
if DATA_PATH.exists():
logger.info(f"Loading cached tasks dataframe from {DATA_PATH}")
return pd.read_parquet(DATA_PATH)
logger.info("Creating tasks dataframe")
query = """
SELECT
tr.onetsoc_code,
tr.task_id,
ts.task,
od.title AS occupation_title,
od.description AS occupation_description,
tr.scale_id,
tr.category,
tr.data_value
FROM
task_ratings tr
JOIN
task_statements ts ON tr.task_id = ts.task_id
JOIN
occupation_data od ON tr.onetsoc_code = od.onetsoc_code;
"""
ratings_df = pd.read_sql_query(query, self.onet_conn)
logger.info(f"Fetched {len(ratings_df)} task rating records from the database.")
# 1. Handle Frequency (FT)
logger.info("Processing Frequency data")
freq_df = ratings_df[ratings_df["scale_id"] == "FT"].copy()
if not freq_df.empty:
freq_pivot = freq_df.pivot_table(
index=["onetsoc_code", "task_id"],
columns="category",
values="data_value",
fill_value=0,
)
freq_pivot.columns = [f"frequency_category_{int(col)}" for col in freq_pivot.columns]
else:
raise ValueError("No frequency data.")
# 2. Handle Importance (IM, IJ)
logger.info("Processing Importance data")
imp_df = ratings_df[ratings_df["scale_id"].isin(["IM", "IJ"])].copy()
if not imp_df.empty:
imp_avg = imp_df.groupby(["onetsoc_code", "task_id"])["data_value"].mean().reset_index()
imp_avg.rename(columns={"data_value": "importance_average"}, inplace=True)
else:
raise ValueError("No importance data.")
# 3. Handle Relevance (RT)
logger.info("Processing Relevance data")
rel_df = ratings_df[ratings_df["scale_id"] == "RT"].copy()
if not rel_df.empty:
rel_avg = rel_df.groupby(["onetsoc_code", "task_id"])["data_value"].mean().reset_index()
rel_avg.rename(columns={"data_value": "relevance_average"}, inplace=True)
else:
raise ValueError("No relevance data.")
# 5. Get Base Task/Occupation Info
logger.info("Extracting base task/occupation info")
base_cols = ["onetsoc_code", "task_id", "task", "occupation_title", "occupation_description"]
base_info = ratings_df[base_cols].drop_duplicates().set_index(["onetsoc_code", "task_id"])
# 6. Merge Processed ONET Data
logger.info("Merging processed ONET data")
final_df = base_info.merge(freq_pivot, left_index=True, right_index=True, how="left")
final_df = final_df.reset_index()
if not imp_avg.empty:
final_df = final_df.merge(imp_avg, on=["onetsoc_code", "task_id"], how="left")
else:
final_df["importance_average"] = np.nan
if not rel_avg.empty:
final_df = final_df.merge(rel_avg, on=["onetsoc_code", "task_id"], how="left")
else:
final_df["relevance_average"] = np.nan
final_df = final_df.replace({np.nan: None})
# 7. Merge with EPOCH remote data
logger.info("Merging with EPOCH remote data")
final_df = pd.merge(final_df, self.epoch_df[['Task', 'Remote']], left_on='task', right_on='Task', how='left')
final_df = final_df.drop('Task', axis=1).rename(columns={'Remote': 'remote_status'})
logger.info(f"Created tasks dataframe with shape {final_df.shape}")
final_df.to_parquet(DATA_PATH)
return final_df
def _check_for_insanity(self, df: pd.DataFrame):
if df['lb_estimate_in_minutes'].isnull().any():
missing_count = df['lb_estimate_in_minutes'].isnull().sum()
raise ValueError(f"Found {missing_count} atomic tasks with missing 'lb_estimate_in_minutes'.")
if df['ub_estimate_in_minutes'].isnull().any():
missing_count = df['ub_estimate_in_minutes'].isnull().sum()
raise ValueError(f"Found {missing_count} atomic tasks with missing 'ub_estimate_in_minutes'.")
valid_estimates = df.dropna(subset=['lb_estimate_in_minutes', 'ub_estimate_in_minutes'])
impossible_bounds = valid_estimates[
(valid_estimates['lb_estimate_in_minutes'] <= 0) |
(valid_estimates['ub_estimate_in_minutes'] <= 0) |
(valid_estimates['lb_estimate_in_minutes'] > valid_estimates['ub_estimate_in_minutes'])
]
if not impossible_bounds.empty:
raise ValueError(f"Found {len(impossible_bounds)} rows with impossible bounds (e.g., lb > ub or value <= 0).")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run the econtai pipeline.")
parser.add_argument("--output-dir", type=str, default="dist/", help="The directory to write output files to.")
parser.add_argument("--bust-estimability", action="store_true", help="Bust the saved task estimability classification (EXPENSIVE)")
parser.add_argument("--bust-estimates", action="store_true", help="Bust the tasks estimates (EXPENSIVE)")
parser.add_argument("--debug", action="store_true", help="Enable debug mode (e.g., process fewer tasks).")
args = parser.parse_args()
Runner(output_dir=args.output_dir, debug=args.debug, bust_estimability=args.bust_estimability, bust_estimates=args.bust_estimates).run()

222
pipeline/utils.py Normal file
View file

@ -0,0 +1,222 @@
import subprocess
import matplotlib.colors as mcolors
import matplotlib as mpl
import seaborn as sns
import tempfile
import litellm
import time
import math
from tqdm import tqdm
from typing import Any, List, Dict
from .logger import logger
OCCUPATION_MAJOR_CODES = {
'11': 'Management',
'13': 'Business & Financial',
'15': 'Computer & Mathematical',
'17': 'Architecture & Engineering',
'19': 'Life, Physical, & Social Science',
'21': 'Community & Social Service',
'23': 'Legal',
'25': 'Education, Training, & Library',
'27': 'Arts, Design, & Media',
'29': 'Healthcare Practitioners',
'31': 'Healthcare Support',
'33': 'Protective Service',
'35': 'Food Preparation & Serving',
'37': 'Building & Grounds Maintenance',
'39': 'Personal Care & Service',
'41': 'Sales & Related',
'43': 'Office & Admin Support',
'45': 'Farming, Fishing, & Forestry',
'47': 'Construction & Extraction',
'49': 'Installation, Maintenance, & Repair',
'51': 'Production',
'53': 'Transportation & Material Moving',
'55': 'Military Specific',
}
GRAY = {'50':'#f8fafc','100':'#f1f5f9','200':'#e2e8f0',
'300':'#cbd5e1','400':'#94a3b8','500':'#64748b',
'600':'#475569','700':'#334155','800':'#1e293b',
'900':'#0f172a','950':'#020617'}
LIME = {'50': '#f7fee7','100': '#ecfcca','200': '#d8f999',
'300': '#bbf451','400': '#9ae600','500': '#83cd00',
'600': '#64a400','700': '#497d00','800': '#3c6300',
'900': '#35530e','950': '#192e03'}
def convert_to_minutes(qty, unit):
"""Converts a quantity in a given unit to minutes."""
return qty * {
"minute": 1,
"hour": 60,
"day": 60 * 24,
"week": 60 * 24 * 7,
"month": 60 * 24 * 30,
"trimester": 60 * 24 * 90,
"semester": 60 * 24 * 180,
"year": 60 * 24 * 365,
}[unit]
def pretty_display(df):
print(df)
return
html_output = df.to_html(index=False)
# Create a temporary HTML file
with tempfile.NamedTemporaryFile(mode='w', suffix=".html", encoding="utf-8") as temp_file:
temp_file.write(html_output)
temp_file_path = temp_file.name
subprocess.run(["/home/felix/.nix-profile/bin/firefox-devedition", "-p", "Work (YouthAI)", temp_file_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
input("Press Enter to continue after reviewing the HTML output...")
def enrich(
model: str,
rpm: int, # Requests per minute
messages_to_process: List[List[Dict[str, str]]],
schema: Dict[str, Any],
chunk_size: int = 100,
):
all_results = []
num_messages = len(messages_to_process)
if num_messages == 0:
return all_results
num_chunks = math.ceil(num_messages / chunk_size)
logger.info(f"Starting enrichment for {num_messages} messages, in {num_chunks} chunks of up to {chunk_size} each.")
# Calculate the time that should be allocated per request to respect the RPM limit.
time_per_request = 60.0 / rpm if rpm > 0 else 0
for i in tqdm(range(num_chunks), desc="Enriching data in chunks"):
chunk_start_time = time.time()
start_index = i * chunk_size
end_index = start_index + chunk_size
message_chunk = messages_to_process[start_index:end_index]
if not message_chunk:
continue
try:
# Send requests for the entire chunk in a batch for better performance.
responses = litellm.batch_completion(
model=model,
messages=message_chunk,
response_format={
"type": "json_schema",
"json_schema": schema,
},
)
# batch_completion returns the response or an exception object for each message.
# We'll replace exceptions with None as expected by the calling functions.
for response in responses:
if isinstance(response, Exception):
logger.error(f"API call within batch failed: {response}")
all_results.append(None)
else:
all_results.append(response)
except Exception as e:
# This catches catastrophic failures in batch_completion itself (e.g., auth)
logger.error(f"litellm.batch_completion call failed for chunk {i+1}/{num_chunks}: {e}")
all_results.extend([None] * len(message_chunk))
chunk_end_time = time.time()
elapsed_time = chunk_end_time - chunk_start_time
# To enforce the rate limit, we calculate how long the chunk *should* have taken
# and sleep for the remainder of that time.
if time_per_request > 0:
expected_duration_for_chunk = len(message_chunk) * time_per_request
if elapsed_time < expected_duration_for_chunk:
sleep_duration = expected_duration_for_chunk - elapsed_time
logger.debug(f"Chunk processed in {elapsed_time:.2f}s. Sleeping for {sleep_duration:.2f}s to respect RPM.")
time.sleep(sleep_duration)
return all_results
def get_contrasting_text_color(bg_color_hex_or_rgba):
if isinstance(bg_color_hex_or_rgba, str):
rgba = mcolors.to_rgba(bg_color_hex_or_rgba)
else:
rgba = bg_color_hex_or_rgba
r, g, b, _ = rgba
luminance = 0.2126 * r + 0.7152 * g + 0.0722 * b
return 'black' if luminance > 0.55 else 'white'
def style_plot():
"""
Applies a consistent and professional style to all plots.
This function sets matplotlib's rcParams for a global effect.
"""
mpl.rcParams.update({
'figure.facecolor': GRAY['50'],
'figure.edgecolor': 'none',
'figure.figsize': (12, 8),
'figure.dpi': 150,
'axes.facecolor': GRAY['50'],
'axes.edgecolor': GRAY['300'],
'axes.grid': True,
'axes.labelcolor': GRAY['800'],
'axes.titlecolor': GRAY['900'],
'axes.titlesize': 18,
'axes.titleweight': 'bold',
'axes.titlepad': 20,
'axes.labelsize': 14,
'axes.labelweight': 'semibold',
'axes.labelpad': 10,
'axes.spines.top': False,
'axes.spines.right': False,
'axes.spines.left': True,
'axes.spines.bottom': True,
'text.color': GRAY['700'],
'xtick.color': GRAY['600'],
'ytick.color': GRAY['600'],
'xtick.labelsize': 12,
'ytick.labelsize': 12,
'xtick.major.size': 0,
'ytick.major.size': 0,
'xtick.minor.size': 0,
'ytick.minor.size': 0,
'xtick.major.pad': 8,
'ytick.major.pad': 8,
'grid.color': GRAY['200'],
'grid.linestyle': '--',
'grid.linewidth': 1,
'legend.frameon': False,
'legend.fontsize': 12,
'legend.title_fontsize': 14,
'legend.facecolor': 'inherit',
'font.family': 'sans-serif',
'font.sans-serif': ['Inter'],
'font.weight': 'normal',
'lines.linewidth': 2,
'lines.markersize': 6,
})
# Seaborn specific styles
# Use shades of LIME as the primary color palette.
# Sorting by integer value of keys, and reversed to have darker shades first.
# Excluding very light colors that won't be visible on a light background.
lime_palette = [LIME[k] for k in sorted(LIME.keys(), key=int, reverse=True) if k not in ['50', '100', '700', '800', '900', '950',]]
sns.set_palette(lime_palette)
sns.set_style("whitegrid", {
'axes.edgecolor': GRAY['300'],
'grid.color': GRAY['200'],
'grid.linestyle': '--',
})

View file

@ -5,28 +5,24 @@ description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"coloraide>=4.6",
"dotenv>=0.9.9",
"graphviz>=0.20.3",
"jupyter>=1.1.1",
"litellm==1.67.0",
"litellm>=1.73.6",
"matplotlib>=3.10.3",
"notebook>=7.4.1",
"openai>=1.76.0",
"openpyxl>=3.1.5",
"pandas>=2.2.3",
"platformdirs>=4.3.8",
"pyarrow>=20.0.0",
"pyvis>=0.3.2",
"requests>=2.32.3",
"scipy",
"pydantic>=2.11.7",
"pytest>=8.4.1",
"python-dotenv>=1.1.1",
"requests>=2.32.4",
"rich>=14.0.0",
"scipy>=1.16.0",
"seaborn>=0.13.2",
"tenacity>=9.1.2",
"tqdm>=4.67.1",
]
[tool.pytest.ini_options]
pythonpath = "src"
pythonpath = "."
addopts = "-v"
asyncio_mode = "auto"

1894
uv.lock generated

File diff suppressed because it is too large Load diff