setup_model.py and setup_optimizer.py moved to seprate files and adde…

[malinjawi]

Description:
This PR addresses issue #225 by refactoring model and optimizer setup functions into setup_model.py and setup_optimizer.py. Key changes include:

Moved setup functions to separate files for better organization.
Added type hints for clarity and improved type checking.

These changes improve code maintainability and readability. Please review and test!

[RobotSail]

Good idea moving out the setup logic. We probably don't need a separate file for each function, so I recommend just moving all of the setup functions to a single file. We already have setup_accelerator.py, so maybe we can move these all there & simply rename the file as setup_objects.py or something similar? (trying to avoid using setup.py)

[RobotSail]

args here is actually not a typing.Any but rather a argparse.Namespace object:

Suggested change

	def setup_optimizer(args: Any, model: torch.nn.Module) -> torch.optim.Optimizer:
	def setup_optimizer(args: argparse.Namespace model: torch.nn.Module) -> torch.optim.Optimizer:

[RobotSail]

Here are the proper types for each arg:

• args: argparse.Namespace
• tokenizer: transformers.PreTrainedTokenizer,
• train_loader: torch.utils.data.DataLoader

Suggested change

	args: Any, tokenizer: Any, train_loader: Any, grad_accum: int
	args: argparse.Namespace, tokenizer: transformers.PreTrainedTokenizer, train_loader: torch.utils.data.DataLoader, grad_accum: int

[RobotSail]

Thanks for this PR Mohammad! I left a few comments, but this looks good so far.

[Maxusmusti]

Hi @malinjawi could you also rebase on the latest main branch when applying review feedback, thanks!