Feat/nllb support

[hspedro]

NLLB Model Support and Architecture Improvements

Overview

This PR introduces support for Facebook's No Language Left Behind (NLLB) translation model alongside the existing M2M100 model. The changes include comprehensive architecture improvements, enhanced documentation, and better model management capabilities.

Key Features

1. NLLB Model Integration

• Added support for NLLB-200 model with multiple size options:
  • Small (600M parameters, ~1.2GB)
  • Large (3.3B parameters, ~6.6GB)
• Implemented automatic language code conversion between ISO and FLORES-200 formats
• Added model-specific optimizations and configurations

2. Model Management Improvements

• Enhanced model download script with support for multiple model types
• Added size-specific model configurations
• Improved model initialization and loading process
• Added backward compatibility for existing model paths

3. API Enhancements

• Added model type selection in translation requests
• Enhanced response format with model information
• Improved error handling for model-specific cases
• Added validation for language code formats

4. Documentation Updates

• Added comprehensive architecture documentation
• Updated API documentation with NLLB-specific details
• Added language code conversion guidelines
• Enhanced deployment documentation

5. Testing Improvements

• Added tests for NLLB model integration
• Enhanced model download tests
• Added language code conversion tests
• Improved test coverage for model-specific features

Technical Details

Model Configuration

MODEL_CONFIGS = {
    "m2m100": {
        "sizes": ["418M", "1.2B", "12B"],
        "default_size": "418M",
        "model_class": "M2M100ForConditionalGeneration",
        "tokenizer_class": "M2M100Tokenizer",
    },
    "nllb": {
        "sizes": ["600M", "3.3B"],
        "default_size": "600M",
        "model_class": "AutoModelForSeq2SeqLM",
        "tokenizer_class": "AutoTokenizer",
    }
}

API Changes

• Added optional model_type parameter to translation requests
• Enhanced response format with model architecture information
• Added support for FLORES-200 language codes

Environment Variables

• Added BABELTRON_MODEL_TYPE for default model selection
• Enhanced model path configuration options
• Added model size configuration

Breaking Changes

• None. The changes are backward compatible with existing M2M100 usage.

Migration Guide

1. Update environment variables if needed:

   BABELTRON_MODEL_TYPE=nllb  # Optional, defaults to m2m100

2. Download NLLB models if desired:

   make download-model-nllb-small
   # or
   make download-model-nllb-large

3. Update language codes when using NLLB:
   • Use FLORES-200 format (e.g., "eng_Latn" instead of "en")
   • ISO codes are automatically converted when possible

Testing

• All tests passing
• Added comprehensive test coverage for new features
• Verified backward compatibility

Documentation

• Updated architecture documentation with new diagrams
• Enhanced API documentation
• Added model-specific guidelines
• Updated deployment instructions

Future Considerations

1. Additional model size options
2. Model fine-tuning capabilities
3. Enhanced language code handling
4. Performance optimizations

Related Issues

• Closes #XXX (Add NLLB model support)
• Addresses #YYY (Improve model management)
• Resolves #ZZZ (Enhance documentation)

Checklist

• Added NLLB model support
• Updated documentation
• Added tests
• Verified backward compatibility
• Updated deployment guides
• Added language code conversion
• Enhanced error handling
• Updated architecture diagrams