switched the model to Salesforce/codet5p-220m

analyzer.py

@@ -211,116 +211,38 @@ logger.addHandler(handler)
 def analyze_code(tokenizer, model, language, code):
     start_time = time.time()
 
-    messages = [
-        {
-            "role": "system",
-            "content": (
-                "You are a helpful and expert-level AI code reviewer and bug fixer. "
-                "Your task is to analyze the given buggy code in the specified programming language, "
-                "identify bugs (logical, syntax, runtime, etc.), and fix them. "
-                "Return a JSON object with the following keys:\n\n"
-                "1. 'bug_analysis': a list of objects, each containing:\n"
-                "   - 'line_number': the line number (approximate if needed)\n"
-                "   - 'error_message': a short name of the bug\n"
-                "   - 'explanation': short explanation of the problem\n"
-                "   - 'fix_suggestion': how to fix it\n"
-                "2. 'corrected_code': the entire corrected code block.\n\n"
-                "Respond only with a JSON block, no extra commentary."
-            )
-        },
-        {
-            "role": "user",
-            "content": f"💻 Language: {language}\n🐞 Buggy Code:\n```{language.lower()}\n{code.strip()}\n```"
-        }
-    ]
+    prompt = (
+        f"Language: {language}\n"
+        f"Task: Fix the following buggy code and explain the bugs.\n"
+        f"Input Code:\n{code.strip()}\n\n"
+        f"Respond with a JSON like this:\n"
+        f"{{\n"
+        f"  \"bug_analysis\": [{{\"line_number\": X, \"error_message\": \"...\", \"explanation\": \"...\", \"fix_suggestion\": \"...\"}}],\n"
+        f"  \"corrected_code\": \"...\"\n"
+        f"}}"
+    )
 
     try:
-        logger.info("📦 Tokenizing input...")
-        inputs = tokenizer.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            return_tensors="pt"
-        ).to(model.device)
-
-        attention_mask = (inputs != tokenizer.pad_token_id).long()
-
-        logger.info("⚙️ Starting generation...")
-        generation_start = time.time()
-        outputs = model.generate(
-            inputs,
-            attention_mask=attention_mask,
-            max_new_tokens=1024,
-            do_sample=False,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id
-        )
-        generation_time = time.time() - generation_start
-        logger.info(f"⚡ Generation completed in {generation_time:.2f} seconds")
-
-        logger.info("📝 Decoding response...")
-        response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-
-        logger.info(f"📄 Response length: {len(response)} characters")
-        logger.info(f"🔍 First 100 chars: {response[:100]}...")
-
-        # Attempt to parse as JSON
-        logger.info("🔍 Attempting to parse JSON...")
-        cleaned_response = response.strip()
-        if cleaned_response.startswith('```json'):
-            cleaned_response = cleaned_response[7:]
-        elif cleaned_response.startswith('```'):
-            cleaned_response = cleaned_response[3:]
-        if cleaned_response.endswith('```'):
-            cleaned_response = cleaned_response[:-3]
-
-        cleaned_response = cleaned_response.strip()
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(model.device)
+        output = model.generate(**inputs, max_new_tokens=1024)
+        response = tokenizer.decode(output[0], skip_special_tokens=True)
 
-        json_output = json.loads(cleaned_response)
+        # Clean response if needed
+        cleaned = response.strip().strip("```json").strip("```").strip()
+        json_output = json.loads(cleaned)
 
-        total_time = time.time() - start_time
-        logger.info(f"✅ Analysis completed successfully in {total_time:.2f} seconds")
-
-        # Validate and patch missing keys
-        if not isinstance(json_output, dict):
-            raise ValueError("Parsed response is not a dictionary")
-
-        if 'bug_analysis' not in json_output:
-            logger.warning("⚠️ Missing 'bug_analysis' key, adding empty list")
-            json_output['bug_analysis'] = []
-
-        if 'corrected_code' not in json_output:
-            logger.warning("⚠️ Missing 'corrected_code' key, adding original code")
-            json_output['corrected_code'] = code
-
-        return json_output
-
-    except json.JSONDecodeError as e:
-        logger.error(f"❌ JSON decode error: {e}")
-        logger.error(f"📄 Raw response: {repr(response)}")
         return {
-            "bug_analysis": [{
-                "line_number": 1,
-                "error_message": "Analysis parsing failed",
-                "explanation": "The AI model returned a response that couldn't be parsed as JSON",
-                "fix_suggestion": "Please try again or check the code format"
-            }],
-            "corrected_code": code,
-            "raw_output": response,
-            "parsing_error": str(e)
+            "bug_analysis": json_output.get("bug_analysis", []),
+            "corrected_code": json_output.get("corrected_code", code)
         }
 
     except Exception as e:
-        total_time = time.time() - start_time
-        logger.error(f"❌ Analysis failed after {total_time:.2f} seconds: {str(e)}")
-        logger.error(f"💥 Exception type: {type(e).__name__}")
         return {
             "bug_analysis": [{
-                "line_number": 1,
-                "error_message": "Analysis failed",
-                "explanation": f"An error occurred during analysis: {str(e)}",
-                "fix_suggestion": "Please try again or contact support"
+                "line_number": 0,
+                "error_message": "Failed to parse",
+                "explanation": str(e),
+                "fix_suggestion": "Try simplifying the code."
             }],
-            "corrected_code": code,
-            "error": str(e),
-            "error_type": type(e).__name__
+            "corrected_code": code
         }

main.py

@@ -295,7 +295,7 @@ async def analyze(req: AnalyzeRequest):
     
     try:
         tokenizer, model = get_model()
-        result = analyze_code(req.language, req.code, tokenizer, model)
+        result = analyze_code(tokenizer, model, req.language, req.code)
         
         if result is None:
             raise HTTPException(status_code=500, detail="Model failed to return any response.")
@@ -350,7 +350,8 @@ async def analyze_for_frontend(req: AnalyzeRequest):
     
     try:
         tokenizer, model = get_model()
-        result = analyze_code(req.language, req.code, tokenizer, model)
+        result = analyze_code(tokenizer, model, req.language, req.code)
+
         
         if result is None:
             return {

model.py

@@ -1,124 +1,159 @@
-# model.py - Optimized version
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+# # model.py - Optimized version
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+# import torch
+# from functools import lru_cache
+# import os
+# import asyncio
+# from concurrent.futures import ThreadPoolExecutor
+# import logging
+
+# logger = logging.getLogger(__name__)
+
+# # Global variables to store loaded model
+# _tokenizer = None
+# _model = None
+# _model_loading = False
+# _model_loaded = False
+
+# @lru_cache(maxsize=1)
+# def get_model_config():
+#     """Cache model configuration"""
+#     return {
+#         "model_id": "deepseek-ai/deepseek-coder-1.3b-instruct",
+#         "torch_dtype": torch.bfloat16,
+#         "device_map": "auto",
+#         "trust_remote_code": True,
+#         # Add these optimizations
+#         "low_cpu_mem_usage": True,
+#         "use_cache": True,
+#     }
+
+# def load_model_sync():
+#     """Synchronous model loading with optimizations"""
+#     global _tokenizer, _model, _model_loaded
+    
+#     if _model_loaded:
+#         return _tokenizer, _model
+    
+#     config = get_model_config()
+#     model_id = config["model_id"]
+    
+#     logger.info(f"🔧 Loading model {model_id}...")
+    
+#     try:
+#         # Set cache directory to avoid re-downloading
+#         cache_dir = os.environ.get("TRANSFORMERS_CACHE", "./model_cache")
+#         os.makedirs(cache_dir, exist_ok=True)
+        
+#         # Load tokenizer first (faster)
+#         logger.info("📝 Loading tokenizer...")
+#         _tokenizer = AutoTokenizer.from_pretrained(
+#             model_id,
+#             trust_remote_code=config["trust_remote_code"],
+#             cache_dir=cache_dir,
+#             use_fast=True,  # Use fast tokenizer if available
+#         )
+        
+#         # Load model with optimizations
+#         logger.info("🧠 Loading model...")
+#         _model = AutoModelForCausalLM.from_pretrained(
+#             model_id,
+#             trust_remote_code=config["trust_remote_code"],
+#             torch_dtype=config["torch_dtype"],
+#             device_map=config["device_map"],
+#             low_cpu_mem_usage=config["low_cpu_mem_usage"],
+#             cache_dir=cache_dir,
+#             offload_folder="offload",  
+#              offload_state_dict=True      
+#         )
+        
+#         # Set to evaluation mode
+#         _model.eval()
+        
+#         _model_loaded = True
+#         logger.info("✅ Model loaded successfully!")
+#         return _tokenizer, _model
+        
+#     except Exception as e:
+#         logger.error(f"❌ Failed to load model: {e}")
+#         raise
+
+# async def load_model_async():
+#     """Asynchronous model loading"""
+#     global _model_loading
+    
+#     if _model_loaded:
+#         return _tokenizer, _model
+    
+#     if _model_loading:
+#         # Wait for ongoing loading to complete
+#         while _model_loading and not _model_loaded:
+#             await asyncio.sleep(0.1)
+#         return _tokenizer, _model
+    
+#     _model_loading = True
+    
+#     try:
+#         # Run model loading in thread pool to avoid blocking
+#         loop = asyncio.get_event_loop()
+#         with ThreadPoolExecutor(max_workers=1) as executor:
+#             tokenizer, model = await loop.run_in_executor(
+#                 executor, load_model_sync
+#             )
+#         return tokenizer, model
+#     finally:
+#         _model_loading = False
+
+# def get_model():
+#     """Get the loaded model (for synchronous access)"""
+#     if not _model_loaded:
+#         return load_model_sync()
+#     return _tokenizer, _model
+
+# def is_model_loaded():
+#     """Check if model is loaded"""
+#     return _model_loaded
+
+# def get_model_info():
+#     """Get model information without loading"""
+#     config = get_model_config()
+#     return {
+#         "model_id": config["model_id"],
+#         "loaded": _model_loaded,
+#         "loading": _model_loading,
+#     }
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from functools import lru_cache
-import os
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
 import logging
 
 logger = logging.getLogger(__name__)
-
-# Global variables to store loaded model
+_model_loaded = False
 _tokenizer = None
 _model = None
-_model_loading = False
-_model_loaded = False
-
 @lru_cache(maxsize=1)
 def get_model_config():
-    """Cache model configuration"""
     return {
-        "model_id": "deepseek-ai/deepseek-coder-1.3b-instruct",
-        "torch_dtype": torch.bfloat16,
-        "device_map": "auto",
-        "trust_remote_code": True,
-        # Add these optimizations
-        "low_cpu_mem_usage": True,
-        "use_cache": True,
+        "model_id": "Salesforce/codet5p-220m",
+        "trust_remote_code": True
     }
 
 def load_model_sync():
-    """Synchronous model loading with optimizations"""
     global _tokenizer, _model, _model_loaded
-    
+
     if _model_loaded:
         return _tokenizer, _model
-    
+
     config = get_model_config()
     model_id = config["model_id"]
-    
-    logger.info(f"🔧 Loading model {model_id}...")
-    
+
     try:
-        # Set cache directory to avoid re-downloading
-        cache_dir = os.environ.get("TRANSFORMERS_CACHE", "./model_cache")
-        os.makedirs(cache_dir, exist_ok=True)
-        
-        # Load tokenizer first (faster)
-        logger.info("📝 Loading tokenizer...")
-        _tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
-            trust_remote_code=config["trust_remote_code"],
-            cache_dir=cache_dir,
-            use_fast=True,  # Use fast tokenizer if available
-        )
-        
-        # Load model with optimizations
-        logger.info("🧠 Loading model...")
-        _model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            trust_remote_code=config["trust_remote_code"],
-            torch_dtype=config["torch_dtype"],
-            device_map=config["device_map"],
-            low_cpu_mem_usage=config["low_cpu_mem_usage"],
-            cache_dir=cache_dir,
-            offload_folder="offload",  
-             offload_state_dict=True      
-        )
-        
-        # Set to evaluation mode
+        _tokenizer = AutoTokenizer.from_pretrained(model_id)
+        _model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
         _model.eval()
-        
         _model_loaded = True
-        logger.info("✅ Model loaded successfully!")
         return _tokenizer, _model
-        
+
     except Exception as e:
         logger.error(f"❌ Failed to load model: {e}")
         raise
-
-async def load_model_async():
-    """Asynchronous model loading"""
-    global _model_loading
-    
-    if _model_loaded:
-        return _tokenizer, _model
-    
-    if _model_loading:
-        # Wait for ongoing loading to complete
-        while _model_loading and not _model_loaded:
-            await asyncio.sleep(0.1)
-        return _tokenizer, _model
-    
-    _model_loading = True
-    
-    try:
-        # Run model loading in thread pool to avoid blocking
-        loop = asyncio.get_event_loop()
-        with ThreadPoolExecutor(max_workers=1) as executor:
-            tokenizer, model = await loop.run_in_executor(
-                executor, load_model_sync
-            )
-        return tokenizer, model
-    finally:
-        _model_loading = False
-
-def get_model():
-    """Get the loaded model (for synchronous access)"""
-    if not _model_loaded:
-        return load_model_sync()
-    return _tokenizer, _model
-
-def is_model_loaded():
-    """Check if model is loaded"""
-    return _model_loaded
-
-def get_model_info():
-    """Get model information without loading"""
-    config = get_model_config()
-    return {
-        "model_id": config["model_id"],
-        "loaded": _model_loaded,
-        "loading": _model_loading,
-    }

setup.py

@@ -1,106 +1,122 @@
-#!/usr/bin/env python3
-"""
-Quick setup script to optimize your existing ML microservice.
-Run this to set up caching and pre-download the model.
-"""
+# #!/usr/bin/env python3
+# """
+# Quick setup script to optimize your existing ML microservice.
+# Run this to set up caching and pre-download the model.
+# """
 
-import os
-import sys
-import logging
-from pathlib import Path
+# import os
+# import sys
+# import logging
+# from pathlib import Path
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# # Configure logging
+# logging.basicConfig(level=logging.INFO)
+# logger = logging.getLogger(__name__)
 
-def setup_cache_directory():
-    """Create cache directory for models"""
-    cache_dir = Path("./model_cache")
-    cache_dir.mkdir(exist_ok=True)
-    logger.info(f"✅ Cache directory created: {cache_dir.absolute()}")
-    return cache_dir
+# def setup_cache_directory():
+#     """Create cache directory for models"""
+#     cache_dir = Path("./model_cache")
+#     cache_dir.mkdir(exist_ok=True)
+#     logger.info(f"✅ Cache directory created: {cache_dir.absolute()}")
+#     return cache_dir
 
-def set_environment_variables():
-    """Set environment variables for optimization"""
-    env_vars = {
-        "TRANSFORMERS_CACHE": "./model_cache",
-        "HF_HOME": "./model_cache", 
-        "TORCH_HOME": "./model_cache",
-        "TOKENIZERS_PARALLELISM": "false",
-        "OMP_NUM_THREADS": "4"
-    }
+# def set_environment_variables():
+#     """Set environment variables for optimization"""
+#     env_vars = {
+#         "TRANSFORMERS_CACHE": "./model_cache",
+#         "HF_HOME": "./model_cache", 
+#         "TORCH_HOME": "./model_cache",
+#         "TOKENIZERS_PARALLELISM": "false",
+#         "OMP_NUM_THREADS": "4"
+#     }
     
-    for key, value in env_vars.items():
-        os.environ[key] = value
-        logger.info(f"Set {key}={value}")
+#     for key, value in env_vars.items():
+#         os.environ[key] = value
+#         logger.info(f"Set {key}={value}")
 
-def pre_download_model():
-    """Pre-download the model to cache"""
-    try:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
+# def pre_download_model():
+#     """Pre-download the model to cache"""
+#     try:
+#         from transformers import AutoTokenizer, AutoModelForCausalLM
         
-        model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"
-        cache_dir = "./model_cache"
+#         model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"
+#         cache_dir = "./model_cache"
         
-        logger.info(f"🔧 Pre-downloading model: {model_id}")
-        logger.info("This may take a few minutes on first run...")
+#         logger.info(f"🔧 Pre-downloading model: {model_id}")
+#         logger.info("This may take a few minutes on first run...")
         
-        # Download tokenizer
-        logger.info("📝 Downloading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
-            cache_dir=cache_dir,
-            trust_remote_code=True
-        )
+#         # Download tokenizer
+#         logger.info("📝 Downloading tokenizer...")
+#         tokenizer = AutoTokenizer.from_pretrained(
+#             model_id,
+#             cache_dir=cache_dir,
+#             trust_remote_code=True
+#         )
         
-        # Download model  
-        logger.info("🧠 Downloading model...")
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            cache_dir=cache_dir,
-            trust_remote_code=True,
-            torch_dtype="auto",  # Let it choose the best dtype
-            low_cpu_mem_usage=True,
-        )
+#         # Download model  
+#         logger.info("🧠 Downloading model...")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             model_id,
+#             cache_dir=cache_dir,
+#             trust_remote_code=True,
+#             torch_dtype="auto",  # Let it choose the best dtype
+#             low_cpu_mem_usage=True,
+#         )
         
-        logger.info("✅ Model downloaded and cached successfully!")
-        logger.info(f"📁 Model cached in: {Path(cache_dir).absolute()}")
+#         logger.info("✅ Model downloaded and cached successfully!")
+#         logger.info(f"📁 Model cached in: {Path(cache_dir).absolute()}")
         
-        # Test that everything works
-        logger.info("🧪 Testing model loading...")
-        del model, tokenizer  # Free memory
+#         # Test that everything works
+#         logger.info("🧪 Testing model loading...")
+#         del model, tokenizer  # Free memory
         
-        return True
+#         return True
         
-    except Exception as e:
-        logger.error(f"❌ Failed to pre-download model: {e}")
-        return False
+#     except Exception as e:
+#         logger.error(f"❌ Failed to pre-download model: {e}")
+#         return False
 
-def main():
-    """Main setup function"""
-    logger.info("🚀 Setting up ML Microservice Optimizations")
-    logger.info("=" * 50)
+# def main():
+#     """Main setup function"""
+#     logger.info("🚀 Setting up ML Microservice Optimizations")
+#     logger.info("=" * 50)
     
-    # Step 1: Setup cache directory
-    setup_cache_directory()
+#     # Step 1: Setup cache directory
+#     setup_cache_directory()
     
-    # Step 2: Set environment variables
-    set_environment_variables()
+#     # Step 2: Set environment variables
+#     set_environment_variables()
     
-    # Step 3: Pre-download model
-    success = pre_download_model()
+#     # Step 3: Pre-download model
+#     success = pre_download_model()
     
-    if success:
-        logger.info("\n✅ Setup completed successfully!")
-        logger.info("📋 Next steps:")
-        logger.info("1. Replace your main.py with the optimized version")
-        logger.info("2. Replace your model.py with the optimized version") 
-        logger.info("3. Run: python main.py")
-        logger.info("\n🚀 Your server will now start much faster!")
-    else:
-        logger.error("\n❌ Setup failed!")
-        logger.error("Please check your internet connection and try again.")
-        sys.exit(1)
+#     if success:
+#         logger.info("\n✅ Setup completed successfully!")
+#         logger.info("📋 Next steps:")
+#         logger.info("1. Replace your main.py with the optimized version")
+#         logger.info("2. Replace your model.py with the optimized version") 
+#         logger.info("3. Run: python main.py")
+#         logger.info("\n🚀 Your server will now start much faster!")
+#     else:
+#         logger.error("\n❌ Setup failed!")
+#         logger.error("Please check your internet connection and try again.")
+#         sys.exit(1)
+
+# if __name__ == "__main__":
+#     main()
+
+# setup.py
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import os
+
+MODEL_ID = "Salesforce/codet5p-220m"
+
+def download_model():
+    print(f"[SETUP] Downloading model: {MODEL_ID}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
+    print("[SETUP] Model and tokenizer downloaded ✅")
 
 if __name__ == "__main__":
-    main()
+    os.makedirs("model_cache", exist_ok=True)
+    download_model()