cognitivess_model/modeling_flax_cognitivess.py

from transformers.models.llama.modeling_flax_llama import (
    FlaxLlamaForCausalLM as FlaxCognitivessForCausalLM,
    FlaxLlamaModel as FlaxCognitivessModel,
)

# You can add more specific code or changes here if needed.