Upload test_model.py

utils/test_model.py

@@ -0,0 +1,743 @@
+"""
+Comprehensive test script for Wildnerve TLM that tests both model functionality and weight loading.
+
+Usage:
+    # Test model inference with custom prompt
+    python test_model.py --prompt "Your test prompt here"
+
+    # Test the weights and maths
+    python test_model.py --check-weights --check-math --diagnostics
+    
+    # Test to verify repos and list weights
+    python test_model.py --verify-repos --list-weights
+    
+    # Test everything
+    python test_model.py --all
+
+    # Test just the weight loading
+    python test_model.py --check-weights
+
+    # Check repository access and list available weights
+    python test_model.py --verify-repos --list-weights
+
+    # Test model inference with custom prompt
+    python test_model.py --prompt "What is quantum computing?"
+"""
+import os
+import sys
+import time
+import logging
+import argparse
+import importlib.util
+from typing import Dict, Any, Optional, List, Tuple
+from pathlib import Path
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+def test_model_loading(prompt: str, verbose: bool = False) -> Dict[str, Any]:
+    """
+    Test if the model loads correctly and can generate responses.
+    
+    Args:
+        prompt: Text prompt to test with
+        verbose: Whether to print detailed diagnostics
+        
+    Returns:
+        Dictionary with test results
+    """
+    results = {
+        "success": False,
+        "model_loaded": False,
+        "response": None,
+        "response_type": None,
+        "elapsed_time": 0,
+        "error": None
+    }
+    
+    try:
+        # Import adapter layer
+        from adapter_layer import WildnerveModelAdapter
+        adapter = WildnerveModelAdapter("")
+        logger.info("Model adapter initialized")
+        
+        # Record start time for performance measurement
+        start_time = time.time()
+        
+        # Try to generate a response
+        logger.info(f"Generating response for: {prompt}")
+        response = adapter.generate(prompt)
+        
+        # Record elapsed time
+        elapsed_time = time.time() - start_time
+        results["elapsed_time"] = elapsed_time
+        results["response"] = response
+        
+        # Check if we got a non-fallback response
+        fallback_phrases = [
+            "I've received your input about",
+            "Processing:",
+            "The model couldn't be properly initialized",
+            "No language model available"
+        ]
+        
+        is_fallback = any(phrase in response for phrase in fallback_phrases)
+        results["response_type"] = "fallback" if is_fallback else "model"
+        results["model_loaded"] = not is_fallback
+        results["success"] = True
+        
+        if verbose:
+            logger.info(f"Response ({len(response)} chars): {response[:100]}...")
+            logger.info(f"Response appears to be from: {'fallback' if is_fallback else 'neural model'}")
+            logger.info(f"Generation took: {elapsed_time:.2f} seconds")
+            
+        return results
+        
+    except Exception as e:
+        logger.error(f"Error testing model: {e}", exc_info=True)
+        results["error"] = str(e)
+        return results
+
+def test_math_capability() -> Dict[str, Any]:
+    """
+    Test the model's math capabilities with various arithmetic expressions.
+    
+    Returns:
+        Dictionary with test results
+    """
+    results = {
+        "success": False,
+        "tests_passed": 0,
+        "tests_total": 0,
+        "details": []
+    }
+    
+    # Test cases: (prompt, expected_contains)
+    math_tests = [
+        ("What is 3 + 4?", "7"),
+        ("What is 12 * 5?", "60"),
+        ("Calculate 18 / 6", "3"),
+        ("What is four multiplied by three?", "12"),
+        ("What is seven plus nine?", "16"),
+        ("Compute 25 - 13", "12")
+    ]
+    
+    try:
+        from adapter_layer import WildnerveModelAdapter
+        adapter = WildnerveModelAdapter("")
+        logger.info("Testing math capabilities...")
+        
+        results["tests_total"] = len(math_tests)
+        
+        for i, (prompt, expected) in enumerate(math_tests):
+            logger.info(f"Math test {i+1}/{len(math_tests)}: {prompt}")
+            
+            try:
+                response = adapter.generate(prompt)
+                passes = expected in response
+                
+                results["details"].append({
+                    "prompt": prompt,
+                    "response": response,
+                    "expected": expected,
+                    "passed": passes
+                })
+                
+                if passes:
+                    results["tests_passed"] += 1
+                    logger.info(f"✓ Test passed: found '{expected}' in response")
+                else:
+                    logger.info(f"✗ Test failed: '{expected}' not found in response")
+                    logger.info(f"Response: {response[:100]}...")
+            except Exception as e:
+                logger.error(f"Error in math test: {e}")
+                results["details"].append({
+                    "prompt": prompt,
+                    "error": str(e),
+                    "passed": False
+                })
+        
+        results["success"] = True
+        return results
+    except Exception as e:
+        logger.error(f"Failed to run math tests: {e}")
+        results["error"] = str(e)
+        return results
+
+def test_weight_loading() -> Dict[str, Any]:
+    """Test loading model weights from local files or HF repository.
+    
+    Returns:
+        Dictionary with test results
+    """
+    results = {
+        "success": False,
+        "local_weights_found": False,
+        "downloaded_weights": False,
+        "weight_files": {},
+        "errors": [],
+        "elapsed_time": 0
+    }
+    
+    try:
+        start_time = time.time()
+        
+        # Try to import load_model_weights
+        try:
+            from load_model_weights import load_model_weights, check_for_local_weights, verify_token
+            
+            # First check token
+            token_verified = verify_token()
+            results["token_verified"] = token_verified
+            
+            # Check for local weights
+            local_weights = check_for_local_weights()
+            results["local_weights_found"] = local_weights
+            
+            if local_weights:
+                results["weight_files"] = {
+                    "transformer": os.environ.get("TLM_TRANSFORMER_WEIGHTS"),
+                    "snn": os.environ.get("TLM_SNN_WEIGHTS")
+                }
+                logger.info("Found local weights")
+            else:
+                # Try downloading weights
+                logger.info("No local weights found, downloading from HF Hub...")
+                weight_files = load_model_weights()
+                
+                if weight_files:
+                    results["downloaded_weights"] = True
+                    results["weight_files"] = weight_files
+                    logger.info(f"Downloaded weights: {list(weight_files.keys())}")
+                else:
+                    logger.warning("Failed to download weights")
+                    results["errors"].append("Failed to download weights")
+        
+        except ImportError as e:
+            logger.error(f"Could not import load_model_weights: {e}")
+            results["errors"].append(f"ImportError: {str(e)}")
+        
+        # Check if we got any weights
+        if results["local_weights_found"] or results["downloaded_weights"]:
+            results["success"] = True
+        
+        # Record elapsed time
+        results["elapsed_time"] = time.time() - start_time
+        
+        return results
+    except Exception as e:
+        logger.error(f"Error testing weight loading: {e}", exc_info=True)
+        results["errors"].append(str(e))
+        results["elapsed_time"] = time.time() - start_time
+        return results
+
+def verify_repositories() -> Dict[str, Any]:
+    """Verify access to model repositories.
+    
+    Returns:
+        Dictionary with verification results
+    """
+    results = {
+        "repositories_checked": 0,
+        "repositories_accessible": 0,
+        "details": {}
+    }
+    
+    try:
+        # Try to import verification function
+        from load_model_weights import verify_repository, verify_token
+        
+        # Get token
+        token = os.environ.get("HF_TOKEN", os.environ.get("HF_API_TOKEN"))
+        token_verified = verify_token()
+        results["token_verified"] = token_verified
+        
+        # First try to get repositories from model_repo_config
+        try:
+            from model_repo_config import get_repo_config
+            config = get_repo_config()
+            repos_to_check = [config.repo_id] + config.alternative_paths
+        except ImportError:
+            # Fallback repositories
+            repos_to_check = [
+                "EvolphTech/Weights",
+                "Wildnerve/tlm-0.05Bx12",
+                "Wildnerve/tlm",
+                "EvolphTech/Checkpoints",
+                "bert-base-uncased"  # Fallback public model
+            ]
+        
+        # Check each repository
+        for repo in repos_to_check:
+            logger.info(f"Verifying repository: {repo}")
+            success, files = verify_repository(repo, token)
+            
+            results["repositories_checked"] += 1
+            if success:
+                results["repositories_accessible"] += 1
+            
+            results["details"][repo] = {
+                "accessible": success,
+                "num_files": len(files) if success else 0,
+                "model_files": [f for f in files if f.endswith('.bin') or f.endswith('.pt')] if success else []
+            }
+        
+        return results
+    except Exception as e:
+        logger.error(f"Error verifying repositories: {e}", exc_info=True)
+        results["error"] = str(e)
+        return results
+
+def list_weight_files() -> Dict[str, Any]:
+    """List available weight files in repositories and locally.
+    
+    Returns:
+        Dictionary with weight file lists
+    """
+    results = {
+        "local_weights": [],
+        "repository_weights": {},
+        "error": None
+    }
+    
+    try:
+        # Check local weight files
+        local_paths = [
+            "/app/Weights/",
+            "./Weights/",
+            "/tmp/hf_cache/",
+            "/tmp/tlm_cache/"
+        ]
+        
+        for base_path in local_paths:
+            if os.path.exists(base_path):
+                for root, _, files in os.walk(base_path):
+                    for file in files:
+                        if file.endswith(('.bin', '.pt', '.pth')):
+                            full_path = os.path.join(root, file)
+                            relative_path = os.path.relpath(full_path, base_path)
+                            results["local_weights"].append({
+                                "path": full_path,
+                                "relative_path": relative_path,
+                                "size_mb": os.path.getsize(full_path) / (1024 * 1024)
+                            })
+        
+        # List repository weight files
+        try:
+            from load_model_weights import list_model_files, verify_token
+            
+            # Get token
+            token = os.environ.get("HF_TOKEN", os.environ.get("HF_API_TOKEN"))
+            token_verified = verify_token()
+            
+            # First try to get repositories from model_repo_config
+            try:
+                from model_repo_config import get_repo_config
+                config = get_repo_config()
+                repos_to_check = [config.repo_id] + config.alternative_paths[:2]  # Only check first few
+            except ImportError:
+                # Fallback repositories
+                repos_to_check = ["EvolphTech/Weights", "Wildnerve/tlm-0.05Bx12"]
+            
+            # Check each repository
+            for repo in repos_to_check:
+                try:
+                    logger.info(f"Listing files in repository: {repo}")
+                    files = list_model_files(repo, token)
+                    results["repository_weights"][repo] = files
+                except Exception as e:
+                    logger.warning(f"Error listing files in {repo}: {e}")
+                    results["repository_weights"][repo] = f"Error: {str(e)}"
+        except ImportError as e:
+            results["error"] = f"Could not import functions to list repository files: {e}"
+        
+        return results
+    except Exception as e:
+        logger.error(f"Error listing weight files: {e}", exc_info=True)
+        results["error"] = str(e)
+        return results
+
+def test_weight_loading_in_model() -> Dict[str, Any]:
+    """Test loading weights into an actual model instance.
+    
+    Returns:
+        Dictionary with test results
+    """
+    results = {
+        "success": False,
+        "model_created": False,
+        "weights_loaded": False,
+        "weight_path": None,
+        "error": None
+    }
+    
+    try:
+        # Try to find or download weights
+        weight_loading_results = test_weight_loading()
+        
+        if not (weight_loading_results["local_weights_found"] or weight_loading_results["downloaded_weights"]):
+            results["error"] = "No weights available to test"
+            return results
+        
+        # Get weight path
+        weight_path = None
+        if "transformer" in weight_loading_results.get("weight_files", {}):
+            weight_path = weight_loading_results["weight_files"]["transformer"]
+        
+        if not weight_path or not os.path.exists(weight_path):
+            results["error"] = f"Weight file not found at {weight_path}"
+            return results
+        
+        results["weight_path"] = weight_path
+        
+        # Try to create a model
+        try:
+            # Try model_Custm first
+            try:
+                import model_Custm
+                if hasattr(model_Custm, "Wildnerve_tlm01"):
+                    logger.info("Creating Wildnerve_tlm01 from model_Custm")
+                    model_class = getattr(model_Custm, "Wildnerve_tlm01")
+                    model = model_class(
+                        vocab_size=50257,  # GPT-2 vocab size
+                        specialization="general",
+                        embedding_dim=768,
+                        num_heads=12,
+                        hidden_dim=768,
+                        num_layers=2,
+                        output_size=50257,
+                        dropout=0.1,
+                        max_seq_length=128
+                    )
+                    results["model_created"] = True
+            except Exception as e:
+                logger.warning(f"Error creating model_Custm: {e}")
+                
+                # Try model_PrTr as fallback
+                try:
+                    import model_PrTr
+                    if hasattr(model_PrTr, "Wildnerve_tlm01"):
+                        logger.info("Creating Wildnerve_tlm01 from model_PrTr")
+                        model_class = getattr(model_PrTr, "Wildnerve_tlm01")
+                        model = model_class(
+                            model_name="gpt2"
+                        )
+                        results["model_created"] = True
+                except Exception as e2:
+                    logger.error(f"Error creating model_PrTr: {e2}")
+                    results["error"] = f"Could not create any model: {e}, {e2}"
+                    return results
+        
+            # Load weights into model
+            if results["model_created"]:
+                from load_model_weights import load_weights_into_model
+                success = load_weights_into_model(model, weight_path, strict=False)
+                results["weights_loaded"] = success
+                
+                if success:
+                    # Try a quick test inference
+                    try:
+                        test_input = "This is a test."
+                        if hasattr(model, "generate"):
+                            output = model.generate(prompt=test_input, max_length=20)
+                            logger.info(f"Test inference output: {output}")
+                            results["test_inference"] = output
+                            results["success"] = True
+                    except Exception as inf_err:
+                        logger.warning(f"Test inference failed: {inf_err}")
+                        # Still mark success if weights loaded
+                        results["success"] = True
+                else:
+                    results["error"] = "Failed to load weights into model"
+        
+        except ImportError as e:
+            results["error"] = f"ImportError: {str(e)}"
+            
+        return results
+    except Exception as e:
+        logger.error(f"Error testing weight loading in model: {e}", exc_info=True)
+        results["error"] = str(e)
+        return results
+
+def run_diagnostics() -> Dict[str, Any]:
+    """
+    Run diagnostics on the model environment and dependencies.
+    
+    Returns:
+        Dictionary with diagnostic results
+    """
+    diagnostics = {
+        "python_version": sys.version,
+        "environment": {},
+        "modules": {},
+        "gpu_available": False,
+        "files_present": {},
+        "model_repo_config": None
+    }
+    
+    # Check environment variables
+    for var in ["MODEL_REPO", "HF_TOKEN", "TLM_TRANSFORMER_WEIGHTS", "TLM_SNN_WEIGHTS", 
+                "LOW_MEMORY_MODE", "CUDA_VISIBLE_DEVICES"]:
+        diagnostics["environment"][var] = os.environ.get(var, "Not set")
+    
+    # Check critical modules
+    for module_name in ["torch", "transformers", "adapter_layer", "model_Custm", "model_PrTr", 
+                       "load_model_weights", "model_repo_config"]:
+        try:
+            module_spec = importlib.util.find_spec(module_name)
+            if module_spec is not None:
+                try:
+                    module = importlib.import_module(module_name)
+                    diagnostics["modules"][module_name] = getattr(module, "__version__", "Available (no version)")
+                except Exception as e:
+                    diagnostics["modules"][module_name] = f"Import error: {e}"
+            else:
+                diagnostics["modules"][module_name] = "Not found"
+        except ImportError:
+            diagnostics["modules"][module_name] = "Not available"
+    
+    # Check for GPU
+    try:
+        import torch
+        diagnostics["gpu_available"] = torch.cuda.is_available()
+        if diagnostics["gpu_available"]:
+            diagnostics["gpu_info"] = torch.cuda.get_device_name(0)
+    except:
+        pass
+    
+    # Check critical files
+    required_files = [
+        "adapter_layer.py",
+        "model_Custm.py",
+        "model_PrTr.py",
+        "model_repo_config.py",
+        "load_model_weights.py",
+        "service_registry.py"
+    ]
+    
+    for filename in required_files:
+        file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), filename)
+        diagnostics["files_present"][filename] = os.path.exists(file_path)
+    
+    # Check model repo config
+    try:
+        from model_repo_config import get_repo_config
+        repo_config = get_repo_config()
+        diagnostics["model_repo_config"] = {
+            "repo_id": repo_config.repo_id,
+            "weight_locations": repo_config.weight_locations[:3] + ["..."],  # First few for brevity
+            "has_auth_token": repo_config.has_auth_token(),
+            "cache_dir": repo_config.cache_dir
+        }
+    except Exception as e:
+        diagnostics["model_repo_config_error"] = str(e)
+    
+    return diagnostics
+
+def main():
+    """Main function to parse arguments and run tests"""
+    parser = argparse.ArgumentParser(description="Comprehensive Wildnerve TLM Model Test Suite")
+    parser.add_argument("--prompt", type=str, default="Tell me about Malaysia's culture", 
+                        help="Prompt text to test (default is non-math to force model loading)")
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
+    parser.add_argument("--check-math", action="store_true", help="Run math capability tests")
+    parser.add_argument("--check-weights", action="store_true", help="Test model weight loading")
+    parser.add_argument("--verify-repos", action="store_true", help="Verify repository access")
+    parser.add_argument("--list-weights", action="store_true", help="List available weight files")
+    parser.add_argument("--test-load", action="store_true", help="Test loading weights into model")
+    parser.add_argument("--diagnostics", action="store_true", help="Run system diagnostics")
+    parser.add_argument("--all", action="store_true", help="Run all tests")
+    parser.add_argument("--output", type=str, help="Save results to JSON file")
+    args = parser.parse_args()
+
+    # If --all specified, enable all tests
+    if args.all:
+        args.check_math = True
+        args.check_weights = True
+        args.verify_repos = True
+        args.list_weights = True
+        args.test_load = True
+        args.diagnostics = True
+
+    # Track overall execution time
+    start_time = time.time()
+    results = {}
+    
+    # Run diagnostics if requested
+    if args.diagnostics:
+        logger.info("Running system diagnostics...")
+        try:
+            diagnostics = run_diagnostics()
+            results["diagnostics"] = diagnostics
+            
+            if args.verbose:
+                logger.info("Diagnostic results:")
+                for category, data in diagnostics.items():
+                    logger.info(f"  {category}:")
+                    if isinstance(data, dict):
+                        for key, value in data.items():
+                            logger.info(f"    - {key}: {value}")
+                    else:
+                        logger.info(f"    - {data}")
+        except Exception as e:
+            logger.error(f"Error in diagnostics: {e}")
+            results["diagnostics_error"] = str(e)
+    
+    # Verify repository access if requested
+    if args.verify_repos:
+        logger.info("Verifying model repository access...")
+        try:
+            repo_results = verify_repositories()
+            results["repository_verification"] = repo_results
+            
+            # Log summary
+            logger.info(f"Repositories checked: {repo_results['repositories_checked']}")
+            logger.info(f"Repositories accessible: {repo_results['repositories_accessible']}")
+            
+            if args.verbose:
+                for repo, details in repo_results["details"].items():
+                    status = "✓" if details["accessible"] else "✗"
+                    logger.info(f"  {status} {repo}: {details['num_files']} files")
+        except Exception as e:
+            logger.error(f"Error verifying repositories: {e}")
+            results["repository_verification_error"] = str(e)
+    
+    # List weight files if requested
+    if args.list_weights:
+        logger.info("Listing available weight files...")
+        try:
+            weight_files = list_weight_files()
+            results["weight_files"] = weight_files
+            
+            # Log summary
+            logger.info(f"Local weight files found: {len(weight_files['local_weights'])}")
+            logger.info(f"Repositories with weights: {len(weight_files['repository_weights'])}")
+            
+            if args.verbose:
+                # Show local weights
+                if weight_files["local_weights"]:
+                    logger.info("Local weight files:")
+                    for weight in weight_files["local_weights"]:
+                        logger.info(f"  - {weight['relative_path']} ({weight['size_mb']:.1f} MB)")
+                
+                # Show repository weights
+                for repo, files in weight_files["repository_weights"].items():
+                    if isinstance(files, list):
+                        logger.info(f"Weights in {repo}: {len(files)} files")
+                        for file in files[:5]:  # Show first 5
+                            logger.info(f"  - {file}")
+                        if len(files) > 5:
+                            logger.info(f"  - ... ({len(files)-5} more)")
+        except Exception as e:
+            logger.error(f"Error listing weight files: {e}")
+            results["weight_files_error"] = str(e)
+    
+    # Test weight loading if requested
+    if args.check_weights:
+        logger.info("Testing model weight loading...")
+        try:
+            weight_loading = test_weight_loading()
+            results["weight_loading"] = weight_loading
+            
+            # Log summary
+            if weight_loading["local_weights_found"]:
+                logger.info("✓ Local weights found")
+                for key, path in weight_loading["weight_files"].items():
+                    if path:
+                        logger.info(f"  - {key}: {path}")
+            elif weight_loading["downloaded_weights"]:
+                logger.info("✓ Weights downloaded successfully")
+                for key, path in weight_loading["weight_files"].items():
+                    if path:
+                        logger.info(f"  - {key}: {path}")
+            else:
+                logger.warning("✗ No weights found or downloaded")
+                if weight_loading["errors"]:
+                    for error in weight_loading["errors"]:
+                        logger.warning(f"  - Error: {error}")
+        except Exception as e:
+            logger.error(f"Error testing weight loading: {e}")
+            results["weight_loading_error"] = str(e)
+    
+    # Test loading weights into model if requested
+    if args.test_load:
+        logger.info("Testing loading weights into model...")
+        try:
+            weight_in_model = test_weight_loading_in_model()
+            results["weight_in_model"] = weight_in_model
+            
+            # Log summary
+            if weight_in_model["success"]:
+                logger.info("✓ Successfully loaded weights into model")
+                logger.info(f"  - Weight path: {weight_in_model['weight_path']}")
+                if "test_inference" in weight_in_model:
+                    logger.info(f"  - Test inference: {weight_in_model['test_inference'][:50]}...")
+            else:
+                logger.warning("✗ Failed to load weights into model")
+                if weight_in_model["error"]:
+                    logger.warning(f"  - Error: {weight_in_model['error']}")
+        except Exception as e:
+            logger.error(f"Error testing weights in model: {e}")
+            results["weight_in_model_error"] = str(e)
+    
+    # Test model loading with the provided prompt
+    logger.info(f"Testing model loading with prompt: {args.prompt}")
+    loading_results = test_model_loading(args.prompt, args.verbose)
+    results["model_loading"] = loading_results
+    
+    # Summary of model loading test
+    if loading_results["success"]:
+        if loading_results["model_loaded"]:
+            logger.info("✅ SUCCESS: Model loaded and generated response")
+            logger.info(f"  - Response: {loading_results['response'][:50]}...")
+            logger.info(f"  - Time: {loading_results['elapsed_time']:.2f} seconds")
+        else:
+            logger.warning("⚠️ PARTIAL: Model adapter works but uses fallback (not neural network)")
+            logger.warning(f"  - Fallback response: {loading_results['response'][:50]}...")
+    else:
+        logger.error("❌ FAILED: Could not load the model")
+        if loading_results["error"]:
+            logger.error(f"  - Error: {loading_results['error']}")
+    
+    # Run math tests if requested
+    if args.check_math:
+        logger.info("Running math capability tests...")
+        math_results = test_math_capability()
+        results["math_tests"] = math_results
+        
+        # Summary of math tests
+        if math_results["success"]:
+            logger.info(f"Math tests: {math_results['tests_passed']}/{math_results['tests_total']} passed")
+            
+            if args.verbose:
+                for i, test in enumerate(math_results["details"]):
+                    status = "✓" if test.get("passed") else "✗"
+                    logger.info(f"  {status} Test {i+1}: {test['prompt']}")
+                    if not test.get("passed"):
+                        logger.info(f"    Expected: {test.get('expected')}")
+                        logger.info(f"    Got: {test.get('response', '')[:50]}...")
+        else:
+            logger.error("Failed to run math tests")
+            if "error" in math_results:
+                logger.error(f"  - Error: {math_results['error']}")
+    
+    # Log total execution time
+    elapsed = time.time() - start_time
+    logger.info(f"All tests completed in {elapsed:.2f} seconds")
+    
+    # Save results if requested
+    if args.output:
+        try:
+            import json
+            with open(args.output, 'w') as f:
+                json.dump(results, f, indent=2)
+            logger.info(f"Results saved to {args.output}")
+        except Exception as e:
+            logger.error(f"Failed to save results: {e}")
+
+if __name__ == "__main__":
+    main()