Upload 2 files

adapter_layer.py

@@ -146,161 +146,41 @@ class Wildnerve_tlm01(nn.Module):
     def generate(self, text_input, max_length=None, **kwargs):
         """Generate text - with lazy model loading"""
         try:
-            # Enhanced math pattern detection
-            import re
-            
-            # Pattern 1: Digit format (4 * 3)
-            digit_pattern = re.compile(r'(\d+)\s*([+\-x*/])\s*(\d+)\s*=?')
-            digit_match = digit_pattern.search(text_input.lower())
-            
-            # Pattern 2: Word format (Four multiplied by three)
-            word_numbers = {
-                'zero': 0, 'one': 1, 'two': 2, 'three': 3, 'four': 4, 'five': 5, 
-                'six': 6, 'seven': 7, 'eight': 8, 'nine': 9, 'ten': 10,
-                'eleven': 11, 'twelve': 12, 'thirteen': 13, 'fourteen': 14, 'fifteen': 15,
-                'sixteen': 16, 'seventeen': 17, 'eighteen': 18, 'nineteen': 19, 'twenty': 20
-            }
-            
-            word_operators = {
-                'plus': '+', 'add': '+', 'added to': '+', 'sum': '+',
-                'minus': '-', 'subtract': '-', 'subtracted from': '-', 'difference': '-',
-                'times': '*', 'multiplied by': '*', 'product': '*', 'multiply': '*',
-                'divided by': '/', 'divide': '/', 'quotient': '/'
-            }
-            
-            word_pattern = re.compile(
-                r'(zero|one|two|three|four|five|six|seven|eight|nine|ten|'
-                r'eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|eighteen|nineteen|twenty)'
-                r'\s+(plus|add|added to|sum|minus|subtract|subtracted from|difference|times|'
-                r'multiplied by|product|multiply|divided by|divide|quotient)\s+'
-                r'(zero|one|two|three|four|five|six|seven|eight|nine|ten|'
-                r'eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|eighteen|nineteen|twenty)',
-                re.IGNORECASE
-            )
-            
-            word_match = word_pattern.search(text_input.lower())
-            
-            # Handle digit format match
-            if digit_match:
-                num1 = int(digit_match.group(1))
-                operator = digit_match.group(2)
-                num2 = int(digit_match.group(3))
-                
-                result = None
-                if operator == '+':
-                    result = num1 + num2
-                elif operator == '-':
-                    result = num1 - num2
-                elif operator in ['x', '*']:
-                    result = num1 * num2
-                elif operator == '/':
-                    result = num1 / num2 if num2 != 0 else "Cannot divide by zero"
-                    
-                if result is not None:
-                    math_response = f"The answer to {num1} {operator} {num2} = {result}"
-                    logger.info(f"Providing direct math answer: {math_response}")
-                    return math_response
-            
-            # Handle word format match
-            elif word_match:
-                num1_word = word_match.group(1).lower()
-                operator_word = word_match.group(2).lower()
-                num2_word = word_match.group(3).lower()
-                
-                num1 = word_numbers.get(num1_word)
-                operator = None
-                for op_word, op_symbol in word_operators.items():
-                    if op_word in operator_word:
-                        operator = op_symbol
-                        break
-                num2 = word_numbers.get(num2_word)
-                
-                if num1 is not None and operator is not None and num2 is not None:
-                    result = None
-                    if operator == '+':
-                        result = num1 + num2
-                    elif operator == '-':
-                        result = num1 - num2
-                    elif operator == '*':
-                        result = num1 * num2
-                    elif operator == '/':
-                        result = num1 / num2 if num2 != 0 else "Cannot divide by zero"
-                        
-                    if result is not None:
-                        math_response = f"The answer to {num1_word} {operator_word} {num2_word} is {result}"
-                        logger.info(f"Providing direct math word answer: {math_response}")
-                        return math_response
-            
-            # Also check for common math question patterns
-            math_question_terms = [
-                "what is", "calculate", "solve", "find", "result of", "answer to",
-                "equals", "equal to"
-            ]
-            
-            if any(term in text_input.lower() for term in math_question_terms):
-                # Simple checks for common patterns like "what is 4 times 3" or "what is four times three"
-                for number_word, number_value in word_numbers.items():
-                    for op_word, op_symbol in word_operators.items():
-                        pattern = f"{number_word}.*{op_word}.*{number_word}"
-                        if re.search(pattern, text_input.lower()):
-                            # Extract the specific numbers and operator
-                            match = re.search(
-                                f"({number_word}).*({op_word}).*({number_word})", 
-                                text_input.lower()
-                            )
-                            if match:
-                                num1 = word_numbers.get(match.group(1))
-                                num2 = word_numbers.get(match.group(3))
-                                operator = op_symbol
-                                
-                                if num1 is not None and num2 is not None:
-                                    result = None
-                                    if operator == '+':
-                                        result = num1 + num2
-                                    elif operator == '-':
-                                        result = num1 - num2
-                                    elif operator == '*':
-                                        result = num1 * num2
-                                    elif operator == '/':
-                                        result = num1 / num2 if num2 != 0 else "Cannot divide by zero"
-                                    
-                                    if result is not None:
-                                        math_response = f"The answer is {result}"
-                                        logger.info(f"Providing direct math answer from question: {math_response}")
-                                        return math_response
-
-            # Special case for "Four multiplied by three"
-            if "four multiplied by three" in text_input.lower():
-                return "The answer is 12"
-            
-            # Try to load model on first use
+            # 1. Load model if not already loaded
             if not self.model_loaded:
+                logger.info("Loading model for first request")
                 self._lazy_load_model()
             
-            # If we have a model now, use it
+            # 2. Let the model handle inference directly with NO pattern matching or rules
             if self.model:
                 try:
-                    logger.info(f"Generating with model: {type(self.model).__name__}")
-                    return self.model.generate(
+                    logger.info(f"Sending prompt directly to neural model: {type(self.model).__name__}")
+                    model_response = self.model.generate(
                         prompt=text_input, 
                         max_length=max_length,
                         **kwargs
                     )
+                    
+                    # Log response for debugging but don't intercept or alter it
+                    logger.info(f"Model generated response of length {len(model_response) if isinstance(model_response, str) else 'unknown'}")
+                    
+                    # Return the raw model response - let the model shine (or fail naturally)
+                    return model_response
+                    
                 except Exception as e:
-                    logger.error(f"Model generation error: {e}")
-                    # Try tokenizer-only response as fallback
+                    # Only log the error but don't substitute with rule-based responses
+                    logger.error(f"Neural model inference error: {e}")
+                    # Continue to basic fallback only if the model completely failed
+            else:
+                logger.warning("No model available - only basic response possible")
             
-            # If we have a tokenizer but no model, use simple responses
-            if self.tokenizer and not self.model:
-                # Try to get a slightly better response with tokenizer
-                tokenized = self.tokenizer(text_input, return_tensors="pt", truncation=True)
-                return f"Processing: {text_input[:50]}..."
-                
-            # If no model or tokenizer, return simple response
-            return f"I've received your input about '{text_input[:30]}...'"
+            # 3. Minimal fallback ONLY if model couldn't be loaded or threw exception
+            if self.tokenizer:
+                return f"The model couldn't be properly initialized. Your input: '{text_input[:30]}...'"
+            return f"No language model available to process: '{text_input[:30]}...'"
             
         except Exception as e:
-            logger.error(f"Error in generate method: {e}")
+            logger.error(f"Critical error in generate method: {e}")
             return f"An error occurred processing your request: {str(e)}"
 
     def _lazy_load_model(self):
@@ -311,28 +191,50 @@ class Wildnerve_tlm01(nn.Module):
             # First initialize tokenizer if not already done
             self._initialize_minimal_tokenizer()
             
-            # Download and load model weights first
-            from load_model_weights import download_model_files, load_weights_into_model, verify_token
-            
-            # First verify token is available
-            token_verified = verify_token()
-            logger.info(f"HF Token verification: {token_verified}")
-            
-            # Get weights from HF repository with more robust error reporting
-            logger.info("Downloading model weights...")
+            # Download and load model weights first with better logging
             try:
-                weight_files = download_model_files(repo_id_base="EvolphTech/Weights")
+                from load_model_weights import download_model_files, load_weights_into_model, verify_token
                 
-                # Add detailed logging about weight files
-                if weight_files:
-                    logger.info(f"Download returned {len(weight_files)} weight files: {list(weight_files.keys())}")
-                else:
-                    logger.warning("No weight files were returned from download_model_files")
+                # First verify token is available
+                token_verified = verify_token()
+                logger.info(f"HF Token verification: {token_verified}")
+                
+                # Get weights from HF repository with more robust error reporting
+                logger.info("Downloading model weights...")
+                try:
+                    # Try multiple repositories in priority order
+                    repositories = [
+                        "EvolphTech/Weights",
+                        "Wildnerve/tlm-0.05Bx12",
+                        "Wildnerve/tlm",
+                        "EvolphTech/Checkpoints"
+                    ]
                     
-            except Exception as e:
-                logger.error(f"Error downloading weights: {str(e)}")
+                    weight_files = None
+                    for repo in repositories:
+                        logger.info(f"Attempting to download weights from {repo}...")
+                        try:
+                            weight_files = download_model_files(repo_id_base=repo)
+                            if weight_files and "transformer" in weight_files:
+                                logger.info(f"Successfully downloaded weights from {repo}")
+                                break
+                        except Exception as repo_error:
+                            logger.warning(f"Failed to download from {repo}: {repo_error}")
+                    
+                    # Add detailed logging about weight files
+                    if weight_files:
+                        logger.info(f"Download returned {len(weight_files)} weight files: {list(weight_files.keys())}")
+                    else:
+                        logger.warning("No weight files were returned from download_model_files")
+                        
+                except Exception as e:
+                    logger.error(f"Error downloading weights: {str(e)}")
+                    weight_files = {}
+            except ImportError:
+                logger.error("Could not import load_model_weights - missing dependencies?")
                 weight_files = {}
-            
+
+            # Rest of model loading code (unchanged)
             # Try to load model_Custm first
             if "model_Custm" in self.available_models:
                 try:

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()