feature/improve-type-annotations-and-error-handling

[sanowl]

Description

This pull request enhances the training script with improved type annotations and more robust error handling. The changes focus on making the code more maintainable, easier to debug, and less prone to runtime errors.

Changes

• Added more consistent type annotations throughout the code
• Implemented more robust error handling and logging, particularly for file operations and network-related tasks
• Refactored the large _train function into smaller, more focused functions
• Added additional error checking with specific exceptions and informative error messages
• Enhanced logging throughout the code for better debugging and monitoring

Testing

• Manually tested the script with various configurations
• Verified that error messages are clear and helpful
• Checked that type annotations are consistent and correct

Additional Notes

• This PR does not change the core functionality of the training process
• The changes are aimed at improving code quality and developer experience