Upload 5 files

core/model_manager.py

@@ -0,0 +1,74 @@
+import os
+import json
+import logging
+from typing import Optional, Dict, Any
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+class ModelManager:
+    def __init__(self):
+        """Initialize the model manager."""
+        self.current_model = None
+        self.current_tokenizer = None
+        self.current_model_name = None
+        self.load_model()
+        
+    def load_model(self, model_name: Optional[str] = None) -> bool:
+        """
+        Load the language model, trying different models in order of preference.
+        
+        Args:
+            model_name: Optional specific model to load
+            
+        Returns:
+            bool: True if any model was loaded successfully
+        """
+        models_to_try = [
+            model_name
+        ] if model_name else [
+            "mistralai/Mistral-7B-Instruct-v0.2",  # Best balance of capability/size
+            "microsoft/phi-2",                      # Fallback
+            "gpt2"                                  # Last resort
+        ]
+        
+        for model_id in models_to_try:
+            try:
+                logger.info(f"Loading {model_id}")
+                self.current_tokenizer = AutoTokenizer.from_pretrained(model_id)
+                self.current_model = AutoModelForCausalLM.from_pretrained(
+                    model_id,
+                    device_map="auto",
+                    torch_dtype=torch.float16,  # Use half precision
+                    load_in_8bit=True
+                )
+                self.current_model_name = model_id
+                self.current_model.eval()
+                logger.info(f"Successfully loaded {model_id}")
+                return True
+            except Exception as e:
+                logger.warning(f"Failed to load {model_id}: {e}")
+                continue
+        
+        return False
+                torch_dtype=getattr(torch, self.config.get('torch_dtype', 'float32'))
+            )
+            
+            self.current_model.eval()
+            self.current_model_name = model_name
+            
+            logger.info(f"Successfully loaded model {model_name}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error loading model {model_name}: {e}")
+            return False
+            
+    def get_current_model(self) -> tuple:
+        """Get currently loaded model and tokenizer."""
+        return self.current_model, self.current_tokenizer
+        
+    def is_model_loaded(self) -> bool:
+        """Check if a model is currently loaded."""
+        return self.current_model is not None and self.current_tokenizer is not None

core/new_model_manager.py

@@ -0,0 +1,60 @@
+import logging
+import torch
+from typing import Optional, Tuple
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+class ModelManager:
+    def __init__(self):
+        """Initialize the model manager."""
+        self.current_model = None
+        self.current_tokenizer = None
+        self.current_model_name = None
+        self.load_model()
+        
+    def load_model(self, model_name: Optional[str] = None) -> bool:
+        """
+        Load the language model, trying different models in order of preference.
+        
+        Args:
+            model_name: Optional specific model to load
+            
+        Returns:
+            bool: True if any model was loaded successfully
+        """
+        models_to_try = [
+            model_name
+        ] if model_name else [
+            "mistralai/Mistral-7B-Instruct-v0.2",  # Best balance of capability/size
+            "microsoft/phi-2",                      # Fallback
+            "gpt2"                                  # Last resort
+        ]
+        
+        for model_id in models_to_try:
+            try:
+                logger.info(f"Loading {model_id}")
+                self.current_tokenizer = AutoTokenizer.from_pretrained(model_id)
+                self.current_model = AutoModelForCausalLM.from_pretrained(
+                    model_id,
+                    device_map="auto",
+                    torch_dtype=torch.float16,  # Use half precision
+                    load_in_8bit=True
+                )
+                self.current_model_name = model_id
+                self.current_model.eval()
+                logger.info(f"Successfully loaded {model_id}")
+                return True
+            except Exception as e:
+                logger.warning(f"Failed to load {model_id}: {e}")
+                continue
+        
+        return False
+        
+    def get_current_model(self) -> Tuple[Optional[AutoModelForCausalLM], Optional[AutoTokenizer]]:
+        """Get currently loaded model and tokenizer."""
+        return self.current_model, self.current_tokenizer
+        
+    def is_model_loaded(self) -> bool:
+        """Check if a model is currently loaded."""
+        return self.current_model is not None and self.current_tokenizer is not None

core/setup.py

@@ -0,0 +1,67 @@
+import os
+import sys
+import json
+import torch
+import logging
+from pathlib import Path
+
+def setup_environment():
+    """Set up the environment for Codette with modern language models."""
+    
+    # Configure logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    )
+    logger = logging.getLogger(__name__)
+    
+    # Create necessary directories
+    config_dir = Path(__file__).parent.parent / 'config'
+    config_dir.mkdir(exist_ok=True)
+    
+    # Check system requirements
+    logger.info("Checking system requirements...")
+    
+    # Check Python version
+    python_version = sys.version_info
+    if python_version.major < 3 or (python_version.major == 3 and python_version.minor < 8):
+        logger.error("Python 3.8 or higher is required")
+        return False
+        
+    # Check CUDA availability
+    cuda_available = torch.cuda.is_available()
+    if cuda_available:
+        gpu_count = torch.cuda.device_count()
+        gpu_name = torch.cuda.get_device_name(0) if gpu_count > 0 else "Unknown"
+        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3 if gpu_count > 0 else 0
+        
+        logger.info(f"CUDA is available with {gpu_count} device(s)")
+        logger.info(f"GPU: {gpu_name} with {gpu_memory:.1f}GB memory")
+    else:
+        logger.warning("CUDA is not available - running in CPU mode will be very slow")
+        
+    # Create or update environment configuration
+    env_config = {
+        'cuda_available': cuda_available,
+        'gpu_count': gpu_count if cuda_available else 0,
+        'gpu_memory': gpu_memory if cuda_available else 0,
+        'python_version': f"{python_version.major}.{python_version.minor}.{python_version.micro}",
+        'torch_version': torch.__version__
+    }
+    
+    with open(config_dir / 'environment.json', 'w') as f:
+        json.dump(env_config, f, indent=2)
+        
+    # Set environment variables
+    os.environ['TRANSFORMERS_CACHE'] = str(Path.home() / '.cache' / 'huggingface')
+    os.environ['TORCH_HOME'] = str(Path.home() / '.cache' / 'torch')
+    
+    # Optimize for inference
+    if cuda_available:
+        torch.backends.cudnn.benchmark = True
+        
+    logger.info("Environment setup complete")
+    return True
+
+if __name__ == '__main__':
+    setup_environment()