code: def resolve_primary_keys_in_schema(sql_tokens: List[str],
                                     schema: Dict[str, List[TableColumn]]) -> List[str]:
      
      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
docs: 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.