Sushyy/bug-explainer-ml

1# #!/usr/bin/env python32# """3# Quick setup script to optimize your existing ML microservice.4# Run this to set up caching and pre-download the model.5# """6 7# import os8# import sys9# import logging10# from pathlib import Path11 12# # Configure logging13# logging.basicConfig(level=logging.INFO)14# logger = logging.getLogger(__name__)15 16# def setup_cache_directory():17#     """Create cache directory for models"""18#     cache_dir = Path("./model_cache")19#     cache_dir.mkdir(exist_ok=True)20#     logger.info(f"✅ Cache directory created: {cache_dir.absolute()}")21#     return cache_dir22 23# def set_environment_variables():24#     """Set environment variables for optimization"""25#     env_vars = {26#         "TRANSFORMERS_CACHE": "./model_cache",27#         "HF_HOME": "./model_cache", 28#         "TORCH_HOME": "./model_cache",29#         "TOKENIZERS_PARALLELISM": "false",30#         "OMP_NUM_THREADS": "4"31#     }32    33#     for key, value in env_vars.items():34#         os.environ[key] = value35#         logger.info(f"Set {key}={value}")36 37# def pre_download_model():38#     """Pre-download the model to cache"""39#     try:40#         from transformers import AutoTokenizer, AutoModelForCausalLM41        42#         model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"43#         cache_dir = "./model_cache"44        45#         logger.info(f"🔧 Pre-downloading model: {model_id}")46#         logger.info("This may take a few minutes on first run...")47        48#         # Download tokenizer49#         logger.info("📝 Downloading tokenizer...")50#         tokenizer = AutoTokenizer.from_pretrained(51#             model_id,52#             cache_dir=cache_dir,53#             trust_remote_code=True54#         )55        56#         # Download model  57#         logger.info("🧠 Downloading model...")58#         model = AutoModelForCausalLM.from_pretrained(59#             model_id,60#             cache_dir=cache_dir,61#             trust_remote_code=True,62#             torch_dtype="auto",  # Let it choose the best dtype63#             low_cpu_mem_usage=True,64#         )65        66#         logger.info("✅ Model downloaded and cached successfully!")67#         logger.info(f"📁 Model cached in: {Path(cache_dir).absolute()}")68        69#         # Test that everything works70#         logger.info("🧪 Testing model loading...")71#         del model, tokenizer  # Free memory72        73#         return True74        75#     except Exception as e:76#         logger.error(f"❌ Failed to pre-download model: {e}")77#         return False78 79# def main():80#     """Main setup function"""81#     logger.info("🚀 Setting up ML Microservice Optimizations")82#     logger.info("=" * 50)83    84#     # Step 1: Setup cache directory85#     setup_cache_directory()86    87#     # Step 2: Set environment variables88#     set_environment_variables()89    90#     # Step 3: Pre-download model91#     success = pre_download_model()92    93#     if success:94#         logger.info("\n✅ Setup completed successfully!")95#         logger.info("📋 Next steps:")96#         logger.info("1. Replace your main.py with the optimized version")97#         logger.info("2. Replace your model.py with the optimized version") 98#         logger.info("3. Run: python main.py")99#         logger.info("\n🚀 Your server will now start much faster!")100#     else:101#         logger.error("\n❌ Setup failed!")102#         logger.error("Please check your internet connection and try again.")103#         sys.exit(1)104 105# if __name__ == "__main__":106#     main()107 108# setup.py109from transformers import AutoTokenizer, AutoModelForSeq2SeqLM110import os111 112MODEL_ID = "Salesforce/codet5p-220m"113 114def download_model():115    print(f"[SETUP] Downloading model: {MODEL_ID}")116    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)117    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)118    print("[SETUP] Model and tokenizer downloaded ✅")119 120if __name__ == "__main__":121    os.makedirs("model_cache", exist_ok=True)122    download_model()123