<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def resolve_primary_keys_in_schema(sql_tokens: List[str], schema: Dict[str, List[TableColumn]]) -> List[str]: """ Some examples in the text2sql datasets use ID as a column reference to the column of a table which has a primary key. This causes problems if you are trying to constrain a grammar to only produce the column names directly, because you don't know what ID refers to. So instead of dealing with that, we just replace it. """

    primary_keys_for_tables = {name: max(columns, key=lambda x: x.is_primary_key).name
                               for name, columns in schema.items()}
    resolved_tokens = []
    for i, token in enumerate(sql_tokens):
        if i > 2:
            table_name = sql_tokens[i - 2]
            if token == "ID" and table_name in primary_keys_for_tables.keys():
                token = primary_keys_for_tables[table_name]
        resolved_tokens.append(token)
    return resolved_tokens