Create setup_helion.py

setup_helion.py

@@ -0,0 +1,516 @@
+"""
+Helion-OSC Easy Setup & Usage Script
+One-file solution for setting up and using Helion-OSC model
+
+This script handles:
+- Automatic dependency installation
+- Model loading from HuggingFace Spaces
+- GPU/CPU detection
+- Memory optimization
+- Simple inference interface
+- Interactive mode
+
+Usage:
+    python setup_helion.py --setup      # First time setup
+    python setup_helion.py --chat       # Interactive chat
+    python setup_helion.py --generate "your prompt here"
+"""
+
+import subprocess
+import sys
+import os
+import logging
+from pathlib import Path
+
+logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
+logger = logging.getLogger(__name__)
+
+
+def install_dependencies():
+    """Install required dependencies"""
+    logger.info("Installing dependencies...")
+    
+    dependencies = [
+        "torch>=2.0.0",
+        "transformers>=4.40.0",
+        "accelerate>=0.25.0",
+        "sentencepiece>=0.1.99",
+        "safetensors>=0.4.0",
+        "bitsandbytes>=0.41.0",
+        "huggingface-hub>=0.19.0"
+    ]
+    
+    for dep in dependencies:
+        logger.info(f"Installing {dep}...")
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", dep, "-q"])
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to install {dep}: {e}")
+    
+    logger.info("✓ Dependencies installed")
+
+
+def check_dependencies():
+    """Check if dependencies are installed"""
+    required = {
+        "torch": "torch",
+        "transformers": "transformers",
+        "accelerate": "accelerate",
+    }
+    
+    missing = []
+    for name, import_name in required.items():
+        try:
+            __import__(import_name)
+        except ImportError:
+            missing.append(name)
+    
+    return missing
+
+
+class HelionOSCEasy:
+    """Easy-to-use wrapper for Helion-OSC model"""
+    
+    def __init__(
+        self,
+        model_name: str = "DeepXR/Helion-OSC",
+        device: str = "auto",
+        use_8bit: bool = False,
+        use_4bit: bool = False,
+        trust_remote_code: bool = True
+    ):
+        """
+        Initialize Helion-OSC with automatic configuration
+        
+        Args:
+            model_name: Model identifier on HuggingFace
+            device: Device to use ("auto", "cuda", "cpu")
+            use_8bit: Use 8-bit quantization (saves memory)
+            use_4bit: Use 4-bit quantization (saves more memory)
+            trust_remote_code: Trust remote code from model
+        """
+        logger.info("="*80)
+        logger.info("HELION-OSC EASY SETUP")
+        logger.info("="*80)
+        
+        # Import here after dependency check
+        import torch
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        
+        self.model_name = model_name
+        self.torch = torch
+        self.AutoTokenizer = AutoTokenizer
+        self.AutoModelForCausalLM = AutoModelForCausalLM
+        
+        # Detect device
+        if device == "auto":
+            if torch.cuda.is_available():
+                self.device = "cuda"
+                logger.info(f"✓ GPU detected: {torch.cuda.get_device_name(0)}")
+                logger.info(f"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+                self.device = "mps"
+                logger.info("✓ Apple Silicon (MPS) detected")
+            else:
+                self.device = "cpu"
+                logger.info("⚠ No GPU detected, using CPU (will be slower)")
+        else:
+            self.device = device
+        
+        # Set quantization
+        self.use_8bit = use_8bit
+        self.use_4bit = use_4bit
+        
+        if self.use_4bit:
+            logger.info("Using 4-bit quantization (lowest memory)")
+        elif self.use_8bit:
+            logger.info("Using 8-bit quantization (reduced memory)")
+        
+        # Check available memory
+        self._check_memory()
+        
+        # Load model
+        logger.info(f"\nLoading model: {model_name}")
+        logger.info("This may take a few minutes on first run...")
+        
+        try:
+            self._load_model(trust_remote_code)
+            logger.info("✓ Model loaded successfully!")
+            self._print_capabilities()
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            logger.info("\nTroubleshooting tips:")
+            logger.info("1. Try with --use-4bit for lower memory usage")
+            logger.info("2. Make sure you have enough RAM/VRAM")
+            logger.info("3. Check internet connection for downloading")
+            raise
+    
+    def _check_memory(self):
+        """Check available memory"""
+        try:
+            import psutil
+            ram_gb = psutil.virtual_memory().total / 1e9
+            ram_available = psutil.virtual_memory().available / 1e9
+            
+            logger.info(f"\nSystem Memory:")
+            logger.info(f"  Total RAM: {ram_gb:.1f} GB")
+            logger.info(f"  Available: {ram_available:.1f} GB")
+            
+            if self.device == "cuda":
+                gpu_mem = self.torch.cuda.get_device_properties(0).total_memory / 1e9
+                logger.info(f"  GPU VRAM: {gpu_mem:.1f} GB")
+                
+                if gpu_mem < 8 and not (self.use_4bit or self.use_8bit):
+                    logger.warning("  ⚠ Low VRAM detected. Consider using --use-4bit")
+            
+            elif ram_available < 16 and not (self.use_4bit or self.use_8bit):
+                logger.warning("  ⚠ Low RAM detected. Consider using --use-4bit")
+        except:
+            pass
+    
+    def _load_model(self, trust_remote_code: bool):
+        """Load tokenizer and model"""
+        # Load tokenizer
+        logger.info("Loading tokenizer...")
+        self.tokenizer = self.AutoTokenizer.from_pretrained(
+            self.model_name,
+            trust_remote_code=trust_remote_code
+        )
+        
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        # Configure model loading
+        model_kwargs = {
+            "trust_remote_code": trust_remote_code,
+            "low_cpu_mem_usage": True
+        }
+        
+        if self.use_4bit:
+            from transformers import BitsAndBytesConfig
+            model_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=self.torch.bfloat16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4"
+            )
+        elif self.use_8bit:
+            model_kwargs["load_in_8bit"] = True
+        else:
+            if self.device == "cuda":
+                model_kwargs["torch_dtype"] = self.torch.bfloat16
+                model_kwargs["device_map"] = "auto"
+            else:
+                model_kwargs["torch_dtype"] = self.torch.float32
+        
+        # Load model
+        logger.info("Loading model weights...")
+        self.model = self.AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            **model_kwargs
+        )
+        
+        if self.device == "cpu" and not (self.use_4bit or self.use_8bit):
+            self.model = self.model.to(self.device)
+        
+        self.model.eval()
+    
+    def _print_capabilities(self):
+        """Print model capabilities"""
+        logger.info("\n" + "="*80)
+        logger.info("MODEL CAPABILITIES")
+        logger.info("="*80)
+        logger.info("✓ Code generation (Python, JavaScript, C++, Java, Rust, Go, etc.)")
+        logger.info("✓ Mathematical reasoning and theorem proving")
+        logger.info("✓ Algorithm design and optimization")
+        logger.info("✓ Code debugging and error fixing")
+        logger.info("✓ Step-by-step problem solving")
+        logger.info("✓ 250K+ token context length")
+        logger.info("="*80)
+    
+    def generate(
+        self,
+        prompt: str,
+        max_length: int = 2048,
+        temperature: float = 0.7,
+        top_p: float = 0.95,
+        top_k: int = 50,
+        do_sample: bool = True,
+        verbose: bool = True
+    ) -> str:
+        """
+        Generate text from prompt
+        
+        Args:
+            prompt: Input prompt
+            max_length: Maximum tokens to generate
+            temperature: Sampling temperature (higher = more creative)
+            top_p: Nucleus sampling parameter
+            top_k: Top-k sampling parameter
+            do_sample: Use sampling (False = greedy)
+            verbose: Print generation info
+            
+        Returns:
+            Generated text
+        """
+        if verbose:
+            logger.info(f"\nGenerating response...")
+            logger.info(f"Prompt length: {len(prompt)} chars")
+        
+        # Tokenize
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        input_length = inputs.input_ids.shape[1]
+        
+        if verbose:
+            logger.info(f"Input tokens: {input_length}")
+        
+        # Generate
+        with self.torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_length=min(max_length, 8192),  # Limit for reasonable speed
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                do_sample=do_sample,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+        
+        # Decode
+        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Remove prompt from output
+        response = generated_text[len(prompt):].strip()
+        
+        if verbose:
+            output_tokens = outputs.shape[1] - input_length
+            logger.info(f"Generated tokens: {output_tokens}")
+        
+        return response
+    
+    def chat(self, system_prompt: str = "You are Helion-OSC, a helpful AI coding assistant."):
+        """Interactive chat mode"""
+        logger.info("\n" + "="*80)
+        logger.info("INTERACTIVE CHAT MODE")
+        logger.info("="*80)
+        logger.info("Commands:")
+        logger.info("  /help     - Show this help")
+        logger.info("  /clear    - Clear conversation")
+        logger.info("  /settings - Change generation settings")
+        logger.info("  /quit     - Exit chat")
+        logger.info("="*80)
+        
+        conversation = []
+        settings = {
+            "temperature": 0.7,
+            "max_length": 2048,
+            "top_p": 0.95
+        }
+        
+        while True:
+            try:
+                user_input = input("\n💬 You: ").strip()
+                
+                if not user_input:
+                    continue
+                
+                if user_input == "/quit":
+                    logger.info("Goodbye!")
+                    break
+                
+                elif user_input == "/help":
+                    logger.info("\nAvailable commands:")
+                    logger.info("  /help     - Show this help")
+                    logger.info("  /clear    - Clear conversation history")
+                    logger.info("  /settings - Adjust generation settings")
+                    logger.info("  /quit     - Exit chat")
+                    continue
+                
+                elif user_input == "/clear":
+                    conversation = []
+                    logger.info("✓ Conversation cleared")
+                    continue
+                
+                elif user_input == "/settings":
+                    logger.info("\nCurrent settings:")
+                    logger.info(f"  Temperature: {settings['temperature']}")
+                    logger.info(f"  Max length: {settings['max_length']}")
+                    logger.info(f"  Top-p: {settings['top_p']}")
+                    
+                    temp = input("New temperature (0.0-2.0, press Enter to skip): ").strip()
+                    if temp:
+                        settings['temperature'] = float(temp)
+                    
+                    max_len = input("New max length (press Enter to skip): ").strip()
+                    if max_len:
+                        settings['max_length'] = int(max_len)
+                    
+                    logger.info("✓ Settings updated")
+                    continue
+                
+                # Build prompt with conversation history
+                conversation.append({"role": "user", "content": user_input})
+                
+                prompt = system_prompt + "\n\n"
+                for msg in conversation:
+                    if msg["role"] == "user":
+                        prompt += f"User: {msg['content']}\n\n"
+                    else:
+                        prompt += f"Assistant: {msg['content']}\n\n"
+                prompt += "Assistant:"
+                
+                # Generate response
+                response = self.generate(
+                    prompt,
+                    max_length=settings['max_length'],
+                    temperature=settings['temperature'],
+                    top_p=settings['top_p'],
+                    verbose=False
+                )
+                
+                conversation.append({"role": "assistant", "content": response})
+                
+                print(f"\n🤖 Helion: {response}")
+                
+            except KeyboardInterrupt:
+                logger.info("\n\nGoodbye!")
+                break
+            except Exception as e:
+                logger.error(f"Error: {e}")
+
+
+def main():
+    """Main CLI interface"""
+    import argparse
+    
+    parser = argparse.ArgumentParser(
+        description="Helion-OSC Easy Setup & Usage",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # First time setup
+  python setup_helion.py --setup
+  
+  # Interactive chat
+  python setup_helion.py --chat
+  
+  # Generate from prompt
+  python setup_helion.py --generate "Write a Python function to sort a list"
+  
+  # Use 4-bit quantization (low memory)
+  python setup_helion.py --chat --use-4bit
+  
+  # Generate with custom settings
+  python setup_helion.py --generate "Solve x^2 = 16" --temperature 0.3 --max-length 1024
+        """
+    )
+    
+    parser.add_argument(
+        "--setup",
+        action="store_true",
+        help="Install dependencies and set up model"
+    )
+    
+    parser.add_argument(
+        "--chat",
+        action="store_true",
+        help="Start interactive chat mode"
+    )
+    
+    parser.add_argument(
+        "--generate",
+        type=str,
+        help="Generate response for a prompt"
+    )
+    
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="DeepXR/Helion-OSC",
+        help="Model name on HuggingFace"
+    )
+    
+    parser.add_argument(
+        "--use-4bit",
+        action="store_true",
+        help="Use 4-bit quantization (lowest memory)"
+    )
+    
+    parser.add_argument(
+        "--use-8bit",
+        action="store_true",
+        help="Use 8-bit quantization"
+    )
+    
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.7,
+        help="Sampling temperature (default: 0.7)"
+    )
+    
+    parser.add_argument(
+        "--max-length",
+        type=int,
+        default=2048,
+        help="Maximum generation length (default: 2048)"
+    )
+    
+    parser.add_argument(
+        "--top-p",
+        type=float,
+        default=0.95,
+        help="Top-p sampling (default: 0.95)"
+    )
+    
+    args = parser.parse_args()
+    
+    # Setup mode
+    if args.setup:
+        logger.info("Setting up Helion-OSC...")
+        install_dependencies()
+        logger.info("\n✓ Setup complete!")
+        logger.info("\nNext steps:")
+        logger.info("  python setup_helion.py --chat")
+        return
+    
+    # Check dependencies
+    missing = check_dependencies()
+    if missing:
+        logger.error(f"Missing dependencies: {', '.join(missing)}")
+        logger.info("Run: python setup_helion.py --setup")
+        return
+    
+    # Initialize model
+    try:
+        helion = HelionOSCEasy(
+            model_name=args.model,
+            use_8bit=args.use_8bit,
+            use_4bit=args.use_4bit
+        )
+    except Exception as e:
+        logger.error(f"Failed to initialize model: {e}")
+        return
+    
+    # Chat mode
+    if args.chat:
+        helion.chat()
+    
+    # Generate mode
+    elif args.generate:
+        response = helion.generate(
+            args.generate,
+            max_length=args.max_length,
+            temperature=args.temperature,
+            top_p=args.top_p
+        )
+        print(f"\n{response}\n")
+    
+    # Default: show help
+    else:
+        logger.info("No action specified. Use --chat or --generate")
+        logger.info("Run with --help for more options")
+
+
+if __name__ == "__main__":
+    main()