Upload 2 files

main.py

@@ -1,59 +1,54 @@
-# Main.py
-import sys
-sys.stdout.reconfigure(line_buffering=True)  # Force line buffering
-import logging
-import sys
-logging.basicConfig(
-    level=logging.DEBUG, 
-    stream=sys.stdout,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    force=True
-)
+# Main.py - Main entry point for Wildnerve-tlm_HF
 import os
 import sys
 import time
-import torch
+import json
+import gc
 import logging
 import argparse
 import importlib
+import threading
 from typing import Dict, Any, Optional, List, Union, Generator, Tuple
+from pathlib import Path
 
-# Add GPU memory monitoring near the top of the file
-try:
-    import torch
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        print(f"GPU Memory: {torch.cuda.memory_allocated()/1e9:.2f}GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.2f}GB")
-except Exception as e:
-    print(f"Error checking GPU memory: {e}")
+# Set up line buffering early
+sys.stdout.reconfigure(line_buffering=True)
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+# Configure logging once at the top level
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    force=True
+)
 logger = logging.getLogger(__name__)
 
-# Add file handler
+# Add file handlers for persistent logs
 file_handler = logging.FileHandler('/tmp/app_debug.log')
 file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
 logger.addHandler(file_handler)
-logger.info("File logging enabled at /tmp/app_debug.log")
 
-# Add a file handler to persist logs
 fh = logging.FileHandler("/tmp/container.log")
 fh.setLevel(logging.DEBUG)
-formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
-fh.setFormatter(formatter)
+fh.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levellevel)s - %(message)s"))
 logging.getLogger().addHandler(fh)
+logger.info("Logging configured")
 
 # Force early initialization of vital environment variables 
 if not os.environ.get("TLM_DATA_DIR"):
     os.environ["TLM_DATA_DIR"] = "/tmp/tlm_data"
 
-# Enforce CPU mode
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-logger.info("Set CUDA_VISIBLE_DEVICES to force CPU mode")
+# Select GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logger.info(f"Using device: {device}")
 
+# Add GPU memory monitoring
 try:
     import torch
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        logger.info(f"GPU Memory: {torch.cuda.memory_allocated()/1e9:.2f}GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.2f}GB")
+        
+    # Apply PyTorch device fixes
     if hasattr(torch, 'set_default_device'):
         torch.set_default_device("cpu")
     _original_device = torch.device
@@ -62,11 +57,11 @@ try:
     torch.device = patched_device
     logger.info("✅ Applied PyTorch device attribute fix")
 except Exception as e:
-    logger.warning(f"Error applying torch fixes: {e}")
+    logger.warning(f"Error with PyTorch setup: {e}")
 
 # Import configuration
 try:
-    from config import app_config, load_config
+    from config import app_config, load_config, get_model_architecture_params
     
     # Create an emergency patch for config.py issue
     if hasattr(app_config, 'TRANSFORMER_CONFIG'):
@@ -100,9 +95,9 @@ except ImportError:
 
 # Import service registry
 try:
-    from service_registry import registry, MODEL, TOKENIZER, MODEL_MANAGER, COMMUNICATOR, PIPELINE
+    from service_registry import registry, MODEL, TOKENIZER, MODEL_MANAGER, COMMUNICATOR, PIPELINE, PRETRAINED_MODEL
 
-    # Import event system types without importing implementation
+    # Import event system types
     from utils.event_system import (
         EVENT_STDP_REQUEST, EVENT_STDP_RESPONSE, EVENT_TOKEN_GENERATED,
         EVENT_USER_INPUT, EVENT_MODEL_REQUEST, EVENT_MODEL_RESPONSE,
@@ -133,6 +128,7 @@ except ImportError as e:
     MODEL_MANAGER = "model_manager"
     COMMUNICATOR = "communicator"
     PIPELINE = "pipeline"
+    PRETRAINED_MODEL = "pretrained_model"  # Added this constant
     
     # Define minimal event constants
     EVENT_STDP_REQUEST = "stdp_request"
@@ -151,7 +147,7 @@ except ImportError as e:
     
     event_bus = EventBus()
 
-from find_weights import find_transformer_weights, find_snn_weights  # Import weight-finding functions
+from find_weights import find_transformer_weights, find_snn_weights
 
 # Import API components
 try:
@@ -169,7 +165,8 @@ except Exception as e:
     def verify_model_repo_access():
         return False
 
-# Add direct config fix function
+# --- Helper functions ---
+
 def fix_config_file(config_path="config.json"):
     """Fix the config file directly"""
     import os
@@ -788,53 +785,6 @@ def setup_environment(args):
         model_dir = os.path.join(data_dir, "models")
         os.makedirs(model_dir, exist_ok=True)
 
-import logging
-from app import app as fastapi_app
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-def main():
-    logger.info("Starting TLM application")
-    # serve FastAPI instead of Gradio
-    import uvicorn
-    uvicorn.run(
-        "app:app",
-        host="0.0.0.0",
-        port=int(os.getenv("PORT",7860)),
-        workers=os.cpu_count() or 1,
-        loop="auto"
-    )
-
-if __name__ == "__main__":
-    import os
-    import uvicorn
-    from app import app
-
-    if __name__ == "__main__":
-        uvicorn.run(
-            app,
-            host="0.0.0.0",
-            port=int(os.getenv("PORT", 7860)),
-            workers=os.cpu_count() or 1
-        )
-
-# Create a main.py that properly initializes everything
-
-import os
-import sys
-import logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Add near the top, after imports but before loading model modules
-try:
-    import generate_tokens_fix
-    logger.info("Applied generate_tokens fix to prevent recursion errors")
-except Exception as e:
-    logger.error(f"Failed to apply generate_tokens fix: {e}")
-
 def initialize_system():
     """Initialize all components in the correct order"""
     logger.info("Starting system initialization")
@@ -852,13 +802,12 @@ def initialize_system():
         tokenizer = TokenizerWrapper(model_name="gpt2")
     
     # Then register tokenizer
-    from service_registry import registry, TOKENIZER, PRETRAINED_MODEL
     registry.register(TOKENIZER, tokenizer, overwrite=True)
     logger.info("Tokenizer registered")
     
     # Initialize pretrained model (GPT-2)
     try:
-        from model_PrTr import Wildnerve_tlm01 as PretrainedModel
+        from model_PrTr import GPT_2 as PretrainedModel
         pretrained = PretrainedModel(model_name="gpt2", tokenizer=tokenizer)
         registry.register(PRETRAINED_MODEL, pretrained, overwrite=True)
         logger.info("GPT-2 pretrained model registered")
@@ -868,24 +817,26 @@ def initialize_system():
     # Now load custom model
     try:
         from model_Custm import Wildnerve_tlm01
+        # Use architecture parameters from config
+        arch_params = get_model_architecture_params()
+        
         model = Wildnerve_tlm01(
-            vocab_size=50257,  # GPT-2 vocab size
+            vocab_size=arch_params["vocab_size"],
             specialization="general",
             dataset_path=None,
             model_name="gpt2",
-            embedding_dim=768,  # Ensure 768-dimensional model
-            num_heads=12,       # 12 heads for 768-dim
-            hidden_dim=768,     # Ensure 768-dimensional model
-            num_layers=12,      # More layers for larger model
-            output_size=50257,  # GPT-2 vocab size
-            dropout=0.1,
-            max_seq_length=1024,  # Increased for 768-dim model
+            embedding_dim=arch_params["embedding_dim"],
+            num_heads=arch_params["num_heads"],  
+            hidden_dim=arch_params["hidden_dim"],
+            num_layers=arch_params["num_layers"],
+            output_size=arch_params["vocab_size"],
+            dropout=arch_params["dropout"],
+            max_seq_length=arch_params["max_seq_length"],
             pooling_mode="last",
             tokenizer=tokenizer
         )
         
         # Register model
-        from service_registry import MODEL
         registry.register(MODEL, model, overwrite=True)
         logger.info("Custom model registered successfully")
         return True
@@ -893,10 +844,24 @@ def initialize_system():
         logger.error(f"Failed to initialize custom model: {e}", exc_info=True)
         return False
 
-# Call initialization at the beginning
-success = initialize_system()
-logger.info(f"Initialization {'successful' if success else 'failed'}")
+def main():
+    """Main application entry point with consolidated functionality"""
+    # Initialize the system first
+    success = initialize_system()
+    logger.info(f"System initialization {'successful' if success else 'failed'}")
+    
+    # Start the server
+    from app import app
+    import uvicorn
+    
+    logger.info("Starting TLM application")
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(os.getenv("PORT", 7860)),
+        workers=os.cpu_count() or 1,
+        loop="auto"
+    )
 
 if __name__ == "__main__":
-    success = initialize_system()
-    logger.info(f"Initialization {'successful' if success else 'failed'}")
+    main()

model_Custm.py

@@ -355,18 +355,18 @@ class Wildnerve_tlm01(nn.Module, AbstractModel):
                 output = output.unsqueeze(1)
                 
             # Apply final projection to vocabulary space
-            output = self.final_layer(output)
+            logits = self.final_layer(output)
             
             # CRITICAL: Ensure output is always 3D [batch_size, seq_length, vocab_size]
-            if output.dim() == 2:
+            if logits.dim() == 2:
                 # If 2D tensor [batch_size, vocab_size], reshape to 3D [batch_size, 1, vocab_size]
-                batch_size, vocab_size = output.shape
+                batch_size, vocab_size = logits.shape
                 logger.info(f"2D tensor: batch_size={batch_size}, vocab_size={vocab_size}")
-                output = output.unsqueeze(1)  # Add sequence dimension
-                logger.info(f"Reshaped 2D output to 3D tensor: {output.shape}")
+                logits = logits.unsqueeze(1)  # Add sequence dimension
+                logger.info(f"Reshaped 2D output to 3D tensor: {logits.shape}")
             
             # Record the output shape and dimensions for debugging
-            logger.info(f"Output shape: {output.shape}, dimensions: {output.dim()}")
+            logger.info(f"Output shape: {logits.shape}, dimensions: {logits.dim()}")
             
             # Calculate loss if labels are provided
             loss = None
@@ -377,16 +377,21 @@ class Wildnerve_tlm01(nn.Module, AbstractModel):
                     logger.info(f"Reshaped labels to {labels.shape}")
                 
                 # Calculate loss with properly shaped tensors
-                batch_size, seq_length, vocab_size = output.shape
+                batch_size, seq_length, vocab_size = logits.shape
                 loss_fct = nn.CrossEntropyLoss()
-                loss = loss_fct(output.reshape(-1, vocab_size), labels)
+                loss = loss_fct(logits.reshape(-1, vocab_size), labels)
                 logger.info(f"Returning loss tensor: {loss.item()}")
             
             # Return the proper format
-            if loss is not None:
-                return loss, output
-            else:
-                return output
+            from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+            return CausalLMOutputWithCrossAttentions(
+                loss=loss,
+                logits=logits,
+                past_key_values=None,
+                hidden_states=None,
+                attentions=None,
+                cross_attentions=None
+            )
                 
         except Exception as e:
             logger.error(f"Error in forward pass: {str(e)}")
@@ -405,7 +410,14 @@ class Wildnerve_tlm01(nn.Module, AbstractModel):
             # CRITICAL: Return a proper 3D tensor even in error case
             dummy_output = torch.zeros((dummy_batch, 1, self.vocab_size), device=next(self.parameters()).device)
             dummy_loss = torch.tensor(float('nan'), device=next(self.parameters()).device)
-            return dummy_loss, dummy_output
+            return CausalLMOutputWithCrossAttentions(
+                loss=dummy_loss,
+                logits=dummy_output,
+                past_key_values=None,
+                hidden_states=None,
+                attentions=None,
+                cross_attentions=None
+            )
 
     # Add sentence transformer methods
     def encode_sentences(self, sentences, batch_size=32, normalize_embeddings=True):