Unable to train drafter for TDM set: KeyError

[wardvl]

I'm trying to train the model for the set TDM.

Preprocessing works fine. I got the csv draft data for the premier draft from 17lands.com and added the expansion in the expansions.py file.

class TDM(Expansion):
    def __init__(
        self,
        bo1=None,
        bo3=None,
        quick=None,
        draft=None,
        replay=None,
        ml_data=False,
        idx_to_name=None,
    ):
        super().__init__(
            expansion="tdm",
            bo1=bo1,
            bo3=bo3,
            quick=quick,
            draft=draft,
            replay=replay,
            ml_data=ml_data,
            idx_to_name=idx_to_name,
        )

    def get_cards_from_scryfall(self):
        tdm = CardSet([f"set=tdm", "is:booster"]) // TDM Booster contents excl basic lands
        special_guests = CardSet([f"set=spg", "cn≥104", "cn≤113"]) // TDM special guests
        all_cards = tdm.union(special_guests)
        tdm.cards = all_cards
        
        basic_lands = CardSet([f"set=tdm", "type:\"basic land\""]) // Basic lands (not included in booster apparently)
        all_cards = tdm.union(basic_lands)
        tdm.cards = all_cards
        return tdm.to_dataframe()

    @property
    def types(self):
        types = super().types
        return types


EXPANSIONS = [VOW, SNC, DMU, BRO, TDM]

I get this error when running the train_drafter.py script:

Traceback (most recent call last):
  File "C:\Users\wardv\mtg\mtg\scripts\train_drafter.py", line 173, in <module>
    main()
  File "C:\Users\wardv\mtg\mtg\scripts\train_drafter.py", line 13, in main  
    train_gen, val_gen = create_train_and_val_gens(
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\mtg\ml\generator.py", line 329, in create_train_and_val_gens
    train_gen = generator(
                ^^^^^^^^^^
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\mtg\ml\generator.py", line 114, in __init__
    super().__init__(
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\mtg\ml\generator.py", line 33, in __init__
    self.generate_global_data(data)
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\mtg\ml\generator.py", line 159, in generate_global_data
    self.pick = data["pick"].apply(lambda x: name_to_idx_mapping[x])        
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^        
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\pandas\core\series.py", line 4924, in apply
    ).apply()
      ^^^^^^^
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\pandas\core\apply.py", line 1427, in apply
    return self.apply_standard()
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\pandas\core\apply.py", line 1507, in apply_standard
    mapped = obj._map_values(
             ^^^^^^^^^^^^^^^^
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\pandas\core\base.py", line 921, in _map_values
    return algorithms.map_array(arr, mapper, na_action=na_action, convert=convert)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\pandas\core\algorithms.py", line 1743, in map_array
    return lib.map_infer(values, mapper, convert=convert)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib.pyx", line 2972, in pandas._libs.lib.map_infer
  File "C:\Users\wardv\mtg\.venv\Lib\site-packages\mtg\ml\generator.py", line 159, in <lambda>
    self.pick = data["pick"].apply(lambda x: name_to_idx_mapping[x])        
                                             ~~~~~~~~~~~~~~~~~~~^^^
KeyError: 'swamp'

I have tried different parameters for the train_drafter.py script. Eg: set exclude_basics=False but nothing seems to work. When I add logging to check what keys are present in the mapping there are no basic lands in there. But the basic lands seem to be in the draft_data.csv file.

Any idea what I'm doing wrong?