DeepXR/Helion-OSC
465
1"""2Helion-OSC Inference Script3DeepXR/Helion-OSC - Mathematical Coding Language Model4 5This module provides comprehensive inference capabilities for the Helion-OSC model,6including specialized methods for different programming and mathematical tasks.7"""8 9import torch10import json11import logging12from typing import Optional, Dict, Any, List, Union13from transformers import (14 AutoTokenizer,15 AutoModelForCausalLM,16 GenerationConfig,17 StoppingCriteria,18 StoppingCriteriaList19)20from dataclasses import dataclass21import warnings22 23# Configure logging24logging.basicConfig(25 level=logging.INFO,26 format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'27)28logger = logging.getLogger(__name__)29 30 31@dataclass32class GenerationParameters:33 """Parameters for text generation"""34 max_length: int = 204835 temperature: float = 0.736 top_p: float = 0.9537 top_k: int = 5038 repetition_penalty: float = 1.0539 length_penalty: float = 1.040 do_sample: bool = True41 num_return_sequences: int = 142 early_stopping: bool = False43 44 45class CodeStoppingCriteria(StoppingCriteria):46 """Custom stopping criteria for code generation"""47 48 def __init__(self, stop_sequences: List[str], tokenizer):49 self.stop_sequences = stop_sequences50 self.tokenizer = tokenizer51 52 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:53 decoded = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)54 return any(seq in decoded for seq in self.stop_sequences)55 56 57class HelionOSCInference:58 """59 Comprehensive inference wrapper for Helion-OSC model60 61 Supports multiple generation modes:62 - Code generation63 - Mathematical reasoning64 - Algorithm design65 - Code debugging66 - Documentation generation67 """68 69 def __init__(70 self,71 model_name: str = "DeepXR/Helion-OSC",72 device: Optional[str] = None,73 load_in_8bit: bool = False,74 load_in_4bit: bool = False,75 use_flash_attention: bool = True,76 trust_remote_code: bool = True77 ):78 """79 Initialize the Helion-OSC model80 81 Args:82 model_name: HuggingFace model identifier83 device: Device to load model on (cuda/cpu/mps)84 load_in_8bit: Load model in 8-bit precision85 load_in_4bit: Load model in 4-bit precision86 use_flash_attention: Use flash attention for faster inference87 trust_remote_code: Trust remote code from model repository88 """89 self.model_name = model_name90 self.device = self._get_device(device)91 self.load_in_8bit = load_in_8bit92 self.load_in_4bit = load_in_4bit93 94 logger.info(f"Initializing Helion-OSC on {self.device}...")95 96 # Load tokenizer97 self.tokenizer = self._load_tokenizer(trust_remote_code)98 99 # Load model100 self.model = self._load_model(101 use_flash_attention=use_flash_attention,102 trust_remote_code=trust_remote_code103 )104 105 # Load generation configs106 self.generation_configs = self._load_generation_configs()107 108 logger.info("Model loaded successfully!")109 self._print_model_info()110 111 def _get_device(self, device: Optional[str]) -> str:112 """Determine the best available device"""113 if device:114 return device115 if torch.cuda.is_available():116 return "cuda"117 elif torch.backends.mps.is_available():118 return "mps"119 return "cpu"120 121 def _load_tokenizer(self, trust_remote_code: bool):122 """Load and configure tokenizer"""123 logger.info("Loading tokenizer...")124 tokenizer = AutoTokenizer.from_pretrained(125 self.model_name,126 trust_remote_code=trust_remote_code,127 padding_side="left"128 )129 130 # Ensure pad token is set131 if tokenizer.pad_token is None:132 tokenizer.pad_token = tokenizer.eos_token133 134 return tokenizer135 136 def _load_model(self, use_flash_attention: bool, trust_remote_code: bool):137 """Load and configure model"""138 logger.info("Loading model...")139 140 model_kwargs = {141 "trust_remote_code": trust_remote_code,142 "low_cpu_mem_usage": True143 }144 145 # Configure precision and quantization146 if self.load_in_8bit:147 model_kwargs["load_in_8bit"] = True148 logger.info("Loading in 8-bit precision")149 elif self.load_in_4bit:150 model_kwargs["load_in_4bit"] = True151 model_kwargs["bnb_4bit_compute_dtype"] = torch.bfloat16152 model_kwargs["bnb_4bit_use_double_quant"] = True153 model_kwargs["bnb_4bit_quant_type"] = "nf4"154 logger.info("Loading in 4-bit precision")155 else:156 if self.device == "cuda":157 model_kwargs["torch_dtype"] = torch.bfloat16158 else:159 model_kwargs["torch_dtype"] = torch.float32160 161 # Configure device mapping162 if self.device == "cuda" and not (self.load_in_8bit or self.load_in_4bit):163 model_kwargs["device_map"] = "auto"164 165 # Load model166 model = AutoModelForCausalLM.from_pretrained(167 self.model_name,168 **model_kwargs169 )170 171 # Move to device if needed172 if self.device != "cuda" or (self.load_in_8bit or self.load_in_4bit):173 if not (self.load_in_8bit or self.load_in_4bit):174 model = model.to(self.device)175 176 model.eval()177 178 # Enable gradient checkpointing for memory efficiency if needed179 if hasattr(model, 'gradient_checkpointing_enable'):180 model.gradient_checkpointing_enable()181 182 return model183 184 def _load_generation_configs(self) -> Dict[str, GenerationParameters]:185 """Load task-specific generation configurations"""186 return {187 "code_generation": GenerationParameters(188 max_length=4096,189 temperature=0.7,190 top_p=0.95,191 top_k=50,192 repetition_penalty=1.05,193 do_sample=True194 ),195 "mathematical_reasoning": GenerationParameters(196 max_length=2048,197 temperature=0.3,198 top_p=0.9,199 top_k=40,200 repetition_penalty=1.0,201 do_sample=False202 ),203 "code_completion": GenerationParameters(204 max_length=1024,205 temperature=0.6,206 top_p=0.92,207 top_k=45,208 repetition_penalty=1.03,209 do_sample=True210 ),211 "algorithm_design": GenerationParameters(212 max_length=3072,213 temperature=0.5,214 top_p=0.93,215 top_k=50,216 repetition_penalty=1.08,217 do_sample=True218 ),219 "debugging": GenerationParameters(220 max_length=2048,221 temperature=0.4,222 top_p=0.88,223 repetition_penalty=1.0,224 do_sample=False225 )226 }227 228 def _print_model_info(self):229 """Print model information"""230 try:231 num_params = sum(p.numel() for p in self.model.parameters())232 logger.info(f"Model parameters: {num_params:,}")233 logger.info(f"Model dtype: {next(self.model.parameters()).dtype}")234 logger.info(f"Device: {self.device}")235 except Exception as e:236 logger.warning(f"Could not get model info: {e}")237 238 def generate(239 self,240 prompt: Union[str, List[str]],241 task_type: str = "code_generation",242 custom_params: Optional[GenerationParameters] = None,243 stop_sequences: Optional[List[str]] = None,244 return_full_text: bool = False,245 **kwargs246 ) -> Union[str, List[str]]:247 """248 Generate text based on prompt249 250 Args:251 prompt: Input prompt or list of prompts252 task_type: Type of task (code_generation, mathematical_reasoning, etc.)253 custom_params: Custom generation parameters254 stop_sequences: List of sequences to stop generation255 return_full_text: Whether to return full text including prompt256 **kwargs: Additional generation parameters257 258 Returns:259 Generated text or list of generated texts260 """261 # Get generation parameters262 if custom_params:263 params = custom_params264 elif task_type in self.generation_configs:265 params = self.generation_configs[task_type]266 else:267 logger.warning(f"Unknown task type '{task_type}', using default parameters")268 params = GenerationParameters()269 270 # Override with kwargs271 for key, value in kwargs.items():272 if hasattr(params, key):273 setattr(params, key, value)274 275 # Tokenize input276 is_batch = isinstance(prompt, list)277 inputs = self.tokenizer(278 prompt,279 return_tensors="pt",280 padding=True,281 truncation=True,282 max_length=self.model.config.max_position_embeddings283 ).to(self.device)284 285 # Setup stopping criteria286 stopping_criteria = None287 if stop_sequences:288 stopping_criteria = StoppingCriteriaList([289 CodeStoppingCriteria(stop_sequences, self.tokenizer)290 ])291 292 # Generate293 with torch.no_grad():294 outputs = self.model.generate(295 **inputs,296 max_length=params.max_length,297 temperature=params.temperature,298 top_p=params.top_p,299 top_k=params.top_k,300 repetition_penalty=params.repetition_penalty,301 length_penalty=params.length_penalty,302 do_sample=params.do_sample,303 num_return_sequences=params.num_return_sequences,304 early_stopping=params.early_stopping,305 pad_token_id=self.tokenizer.pad_token_id,306 eos_token_id=self.tokenizer.eos_token_id,307 stopping_criteria=stopping_criteria308 )309 310 # Decode outputs311 generated_texts = []312 for output in outputs:313 text = self.tokenizer.decode(output, skip_special_tokens=True)314 if not return_full_text and not is_batch:315 # Remove prompt from single generation316 if isinstance(prompt, str):317 text = text[len(prompt):].strip()318 generated_texts.append(text)319 320 return generated_texts if is_batch or params.num_return_sequences > 1 else generated_texts[0]321 322 def code_generation(323 self,324 prompt: str,325 language: Optional[str] = None,326 max_length: int = 4096,327 **kwargs328 ) -> str:329 """330 Generate code for a given prompt331 332 Args:333 prompt: Code generation prompt334 language: Programming language (optional)335 max_length: Maximum length of generated code336 **kwargs: Additional generation parameters337 338 Returns:339 Generated code340 """341 if language:342 prompt = f"Language: {language}\n{prompt}"343 344 return self.generate(345 prompt,346 task_type="code_generation",347 max_length=max_length,348 **kwargs349 )350 351 def mathematical_reasoning(352 self,353 prompt: str,354 max_length: int = 2048,355 **kwargs356 ) -> str:357 """358 Solve mathematical problems with step-by-step reasoning359 360 Args:361 prompt: Mathematical problem362 max_length: Maximum length of solution363 **kwargs: Additional generation parameters364 365 Returns:366 Mathematical solution with reasoning367 """368 return self.generate(369 prompt,370 task_type="mathematical_reasoning",371 max_length=max_length,372 **kwargs373 )374 375 def algorithm_design(376 self,377 prompt: str,378 include_complexity: bool = True,379 max_length: int = 3072,380 **kwargs381 ) -> str:382 """383 Design algorithms with complexity analysis384 385 Args:386 prompt: Algorithm design prompt387 include_complexity: Whether to include complexity analysis388 max_length: Maximum length of output389 **kwargs: Additional generation parameters390 391 Returns:392 Algorithm design with analysis393 """394 if include_complexity:395 prompt += "\n\nPlease include time and space complexity analysis."396 397 return self.generate(398 prompt,399 task_type="algorithm_design",400 max_length=max_length,401 **kwargs402 )403 404 def debug_code(405 self,406 code: str,407 error_message: Optional[str] = None,408 max_length: int = 2048,409 **kwargs410 ) -> str:411 """412 Debug code and provide fixes413 414 Args:415 code: Code to debug416 error_message: Optional error message417 max_length: Maximum length of output418 **kwargs: Additional generation parameters419 420 Returns:421 Debugging analysis and fixes422 """423 prompt = f"Debug the following code:\n\n```\n{code}\n```"424 if error_message:425 prompt += f"\n\nError message: {error_message}"426 prompt += "\n\nProvide a detailed explanation and fixed code."427 428 return self.generate(429 prompt,430 task_type="debugging",431 max_length=max_length,432 **kwargs433 )434 435 def complete_code(436 self,437 code_context: str,438 max_length: int = 1024,439 **kwargs440 ) -> str:441 """442 Complete partial code443 444 Args:445 code_context: Partial code to complete446 max_length: Maximum length of completion447 **kwargs: Additional generation parameters448 449 Returns:450 Code completion451 """452 return self.generate(453 code_context,454 task_type="code_completion",455 max_length=max_length,456 stop_sequences=["\n\n", "```", "###"],457 **kwargs458 )459 460 def batch_generate(461 self,462 prompts: List[str],463 task_type: str = "code_generation",464 batch_size: int = 4,465 **kwargs466 ) -> List[str]:467 """468 Generate responses for multiple prompts in batches469 470 Args:471 prompts: List of prompts472 task_type: Type of task473 batch_size: Batch size for processing474 **kwargs: Additional generation parameters475 476 Returns:477 List of generated responses478 """479 results = []480 for i in range(0, len(prompts), batch_size):481 batch = prompts[i:i + batch_size]482 batch_results = self.generate(batch, task_type=task_type, **kwargs)483 if isinstance(batch_results, str):484 batch_results = [batch_results]485 results.extend(batch_results)486 return results487 488 489def main():490 """Example usage and demonstrations"""491 print("=" * 80)492 print("Helion-OSC Inference Examples")493 print("=" * 80)494 495 # Initialize model496 helion = HelionOSCInference(497 load_in_8bit=False, # Set to True for lower memory usage498 load_in_4bit=False # Set to True for even lower memory usage499 )500 501 # Example 1: Code Generation502 print("\n" + "=" * 80)503 print("Example 1: Code Generation")504 print("=" * 80)505 code_prompt = """Write a Python function to implement a binary search tree with the following methods:506- insert(value): Insert a new value507- search(value): Search for a value508- delete(value): Delete a value509- inorder_traversal(): Return inorder traversal510 511Include proper documentation and type hints."""512 513 print(f"\nPrompt:\n{code_prompt}")514 print("\nGenerating...")515 result = helion.code_generation(code_prompt, language="python")516 print(f"\nGenerated Code:\n{result}")517 518 # Example 2: Mathematical Reasoning519 print("\n" + "=" * 80)520 print("Example 2: Mathematical Reasoning")521 print("=" * 80)522 math_prompt = """Prove that the sum of the first n natural numbers equals n(n+1)/2 using mathematical induction."""523 524 print(f"\nPrompt:\n{math_prompt}")525 print("\nGenerating...")526 result = helion.mathematical_reasoning(math_prompt)527 print(f"\nSolution:\n{result}")528 529 # Example 3: Algorithm Design530 print("\n" + "=" * 80)531 print("Example 3: Algorithm Design")532 print("=" * 80)533 algo_prompt = """Design an efficient algorithm to find the longest palindromic substring in a given string."""534 535 print(f"\nPrompt:\n{algo_prompt}")536 print("\nGenerating...")537 result = helion.algorithm_design(algo_prompt, include_complexity=True)538 print(f"\nAlgorithm:\n{result}")539 540 # Example 4: Code Debugging541 print("\n" + "=" * 80)542 print("Example 4: Code Debugging")543 print("=" * 80)544 buggy_code = """545def fibonacci(n):546 if n <= 1:547 return n548 return fibonacci(n-1) + fibonacci(n-2)549 550# This is too slow for large n551result = fibonacci(100)552"""553 554 print(f"\nBuggy Code:\n{buggy_code}")555 print("\nGenerating debugging analysis...")556 result = helion.debug_code(buggy_code, error_message="Takes too long to compute")557 print(f"\nDebug Analysis:\n{result}")558 559 # Example 5: Batch Processing560 print("\n" + "=" * 80)561 print("Example 5: Batch Code Generation")562 print("=" * 80)563 batch_prompts = [564 "Write a Python function to reverse a linked list",565 "Write a JavaScript function to debounce API calls",566 "Write a Rust function to parse JSON safely"567 ]568 569 print("\nProcessing batch prompts...")570 results = helion.batch_generate(batch_prompts, batch_size=2)571 for i, (prompt, result) in enumerate(zip(batch_prompts, results), 1):572 print(f"\nPrompt {i}: {prompt}")573 print(f"Result {i}:\n{result}\n")574 575 print("=" * 80)576 print("Examples completed!")577 print("=" * 80)578 579 580if __name__ == "__main__":581 main()