File size: 3,355 Bytes
db8cd85 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
#!/usr/bin/env python3
"""
Test script to verify enhanced fallback mechanisms for pre-quantized models.
This simulates the production deployment scenario where bitsandbytes package metadata is missing.
"""
import sys
import logging
import os
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def test_pre_quantized_model_fallback():
"""Test loading a pre-quantized model without bitsandbytes package metadata."""
logger.info("🧪 Testing enhanced fallback for pre-quantized models...")
# Set the problematic model as environment variable
os.environ["AI_MODEL"] = "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit"
try:
from backend_service import current_model, get_quantization_config
from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
logger.info(f"📝 Testing model: {current_model}")
# Test quantization detection
quant_config = get_quantization_config(current_model)
if quant_config:
logger.info(f"✅ Quantization config detected: {type(quant_config).__name__}")
else:
logger.info("📝 No quantization config (bitsandbytes not available)")
# Test the enhanced fallback mechanism
logger.info("🔧 Testing enhanced config-based fallback...")
try:
# This simulates what happens in the lifespan function
config = AutoConfig.from_pretrained(current_model, trust_remote_code=True)
logger.info(f"✅ Successfully loaded config: {type(config).__name__}")
# Check for quantization config in the model config
if hasattr(config, 'quantization_config'):
logger.info(f"🔍 Found quantization_config in model config: {config.quantization_config}")
# Remove it to prevent bitsandbytes errors
config.quantization_config = None
logger.info("🚫 Removed quantization_config from model config")
else:
logger.info("📝 No quantization_config found in model config")
# Test tokenizer loading
logger.info("📥 Testing tokenizer loading...")
tokenizer = AutoTokenizer.from_pretrained(current_model)
logger.info(f"✅ Tokenizer loaded successfully: {len(tokenizer)} tokens")
# Note: We won't actually load the full model in the test to save time/memory
logger.info("✅ Enhanced fallback mechanism validated successfully!")
return True
except Exception as e:
logger.error(f"❌ Enhanced fallback test failed: {e}")
return False
except Exception as e:
logger.error(f"❌ Test setup failed: {e}")
return False
if __name__ == "__main__":
logger.info("🚀 Starting enhanced fallback mechanism test...")
success = test_pre_quantized_model_fallback()
if success:
logger.info("\n🎉 Enhanced fallback test passed!")
logger.info("💡 The deployment should now handle pre-quantized models correctly")
else:
logger.error("\n❌ Enhanced fallback test failed")
sys.exit(0 if success else 1)
|