Add chronos2 encoder as another encoder option in OpenTSLMFlamingo

[liu-jc]

Add chronos2 encoder in OpenTSLMFlamingo

♻️ Current situation & Problem

The previous encoder only use CNNTokenizer, which I believe that it's patchTST style encoder and training from scratch. This PR mainly tried to add chronos2 as the encoder and we can loaded the pretrained chronos2 to utilize their ability to better understand time series.

⚙️ Release Notes

• Chronos2Encoder in src/opentslm/model/encoder/
• Add some wandb funcationalities in curricum_learning.py to faciliate the training log

📚 Documentation

I use in-line comments for describing the implementation, which should be clear.

✅ Testing

Test file in test/test_chronos2_encoder.py

Code of Conduct & Contributing Guidelines

By creating and submitting this pull request, you agree to follow our Code of Conduct and Contributing Guidelines:

• I agree to follow the Code of Conduct and Contributing Guidelines.

[liu-jc]

Hi team @RealLast @max-rosenblattl @masquare ,
As I talked to Patrick, this implementation is adding the pretrained chronos as the encoder in OpenTSLM. The preliminary result shows good.

[RealLast]

Thanks @liu-jc for adding this! Code looks great, just have two minor comments we may want to think about @ThomasKaar @masquare:

1. Not sure if we should add wandb directly in curriculum_learning; we might want to have a helper class or function to enable it externally, to not make curriculum_learning even longer.
2. We should find a better way to modularize encoders, to avoid having to parse a statement like if encoder == "chronos2" for all possible encoders. We could just pass an encoder via the initializer, e.g. model = OpenTSLMFlamingo(.. encoder=Chronos2()).

[liu-jc]

Thanks @liu-jc for adding this! Code looks great, just have two minor comments we may want to think about @ThomasKaar @masquare:

1. Not sure if we should add wandb directly in curriculum_learning; we might want to have a helper class or function to enable it externally, to not make curriculum_learning even longer.
2. We should find a better way to modularize encoders, to avoid having to parse a statement like if encoder == "chronos2" for all possible encoders. We could just pass an encoder via the initializer, e.g. model = OpenTSLMFlamingo(.. encoder=Chronos2()).

Thanks @RealLast for the comments. As we discussed, I will focus more on modularize wandb with a helper class and avoid making curriculum_learning.py even longer.

For modularizing and adding encoders, we maybe delay to another PR later. For this one, we just simply add chronos2encoder. I think the current structure is not too bad. We have the class - "encoder/Chronos2Encoder.py
". which is inherited from TimeSeriesEncoderBase. The only bad thing is that we have to use some statements like encoder == "chronos2"during inference. For this, I think probably we need to modify the inference code and align with the modularization part.

[ThomasKaar]

@liu-jc thanks a lot for the amazing work here!!

given also the comment by @RealLast i would propose a slight different approach that keeps the "original" functionality of the paper intact:

let's create 2 new classes.

• CurriculumTrainierPretrainedEncoder: inherits from CurriculumTrainer and just extends the constructor
• OpenTSLMFlamingoPretrainedEncoder: inherits from OpenTSLMFlamingo and extends the constructor (potentially also just by hardcoding chronos for now but here we could change the encoder pretty dynamically)

and add a new padding functionality to pad_and_apply_batch so we can just naively can call pad_and_apply_batch(..., padding = "left") and _initialize_model` add a new OpenTSLMFlamingoPretrainedEncoder-logic.

that way we might be able to keep additional effort for this PR really low but still make sure it is fully backwards compatible and extends the current framework.

wdyt @RealLast @liu-jc @masquare @max-rosenblattl

[ThomasKaar]

Suggested change

	# print(type(model.vision_encoder), model.vision_encoder)
	# print(dir(model.vision_encoder))

[ThomasKaar]

iiuc this will now change the behavior to chronos by default. this is not ideal since we will need change the behavior of curicullum_training.py and kind of hides the fact that we just chronos.

i would propose another solution: see comment above.

[liu-jc]

Hi @ThomasKaar,
Thanks for the comment! I agree that we may not want to change the default behavior of curicullum_training.py. I can change this back.
But for CurriculumTrainierPretrainedEncoder, I feel that would be somewhat similar to the current implementation just another layer.
For adding a new padding functionality, I agree that it would be great to add padding arg.
Will work on this after ICML ddl.

[masquare]

Thank you for the great work @liu-jc! Please see my comments below.

[masquare]

Can we add an .env.example file to the repo, so users know which environment variables the script expects? and then please add .env to the .gitignore.

[masquare]

please remove personal paths

[masquare]

if you update the dependencies, can you please also update the lockfile using uv lock or uv sync

[liu-jc]

Hi @masquare @ThomasKaar @RealLast,

I updated the code. Basically, removed the personal path, example for .env to support wandb, and added the padding_side.