Create inference.py

inference.py

@@ -0,0 +1,178 @@
+"""
+Helion-V1 Inference Script
+Safe and helpful conversational AI model
+"""
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from typing import List, Dict
+import warnings
+
+warnings.filterwarnings('ignore')
+
+
+class HelionInference:
+    def __init__(self, model_name: str = "DeepXR/Helion-V1", device: str = "auto"):
+        """
+        Initialize the Helion model for inference.
+        
+        Args:
+            model_name: HuggingFace model identifier
+            device: Device to run inference on ('cuda', 'cpu', or 'auto')
+        """
+        print(f"Loading Helion-V1 model from {model_name}...")
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            trust_remote_code=True
+        )
+        
+        self.model.eval()
+        print("Model loaded successfully!")
+        
+        # Safety keywords to monitor
+        self.safety_keywords = [
+            "harm", "illegal", "weapon", "violence", "dangerous",
+            "exploit", "hack", "steal", "abuse"
+        ]
+    
+    def check_safety(self, text: str) -> bool:
+        """
+        Basic safety check on input text.
+        
+        Args:
+            text: Input text to check
+            
+        Returns:
+            True if text appears safe, False otherwise
+        """
+        text_lower = text.lower()
+        for keyword in self.safety_keywords:
+            if keyword in text_lower:
+                return False
+        return True
+    
+    def generate_response(
+        self,
+        messages: List[Dict[str, str]],
+        max_new_tokens: int = 512,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        do_sample: bool = True
+    ) -> str:
+        """
+        Generate a response from the model.
+        
+        Args:
+            messages: List of message dictionaries with 'role' and 'content'
+            max_new_tokens: Maximum number of tokens to generate
+            temperature: Sampling temperature
+            top_p: Nucleus sampling parameter
+            do_sample: Whether to use sampling
+            
+        Returns:
+            Generated response text
+        """
+        # Apply chat template
+        input_ids = self.tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            return_tensors="pt"
+        ).to(self.model.device)
+        
+        # Generate response
+        with torch.no_grad():
+            output = self.model.generate(
+                input_ids,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=do_sample,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id
+            )
+        
+        # Decode response
+        response = self.tokenizer.decode(
+            output[0][input_ids.shape[1]:],
+            skip_special_tokens=True
+        )
+        
+        return response.strip()
+    
+    def chat(self):
+        """Interactive chat mode."""
+        print("\n" + "="*60)
+        print("Helion-V1 Interactive Chat")
+        print("Type 'quit' or 'exit' to end the conversation")
+        print("="*60 + "\n")
+        
+        conversation_history = []
+        
+        while True:
+            user_input = input("You: ").strip()
+            
+            if user_input.lower() in ['quit', 'exit']:
+                print("Goodbye! Have a great day!")
+                break
+            
+            if not user_input:
+                continue
+            
+            # Basic safety check
+            if not self.check_safety(user_input):
+                print("Helion: I apologize, but I can't assist with that request. "
+                      "Let me know if there's something else I can help you with!")
+                continue
+            
+            # Add user message to history
+            conversation_history.append({
+                "role": "user",
+                "content": user_input
+            })
+            
+            # Generate response
+            try:
+                response = self.generate_response(conversation_history)
+                print(f"Helion: {response}\n")
+                
+                # Add assistant response to history
+                conversation_history.append({
+                    "role": "assistant",
+                    "content": response
+                })
+            except Exception as e:
+                print(f"Error generating response: {e}")
+                conversation_history.pop()  # Remove failed user message
+
+
+def main():
+    """Main function for CLI usage."""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Helion-V1 Inference")
+    parser.add_argument("--model", default="DeepXR/Helion-V1", help="Model name or path")
+    parser.add_argument("--device", default="auto", help="Device to use (cuda/cpu/auto)")
+    parser.add_argument("--interactive", action="store_true", help="Start interactive chat")
+    parser.add_argument("--prompt", type=str, help="Single prompt to process")
+    
+    args = parser.parse_args()
+    
+    # Initialize model
+    helion = HelionInference(model_name=args.model, device=args.device)
+    
+    if args.interactive:
+        helion.chat()
+    elif args.prompt:
+        messages = [{"role": "user", "content": args.prompt}]
+        response = helion.generate_response(messages)
+        print(f"Response: {response}")
+    else:
+        print("Please specify --interactive or --prompt")
+
+
+if __name__ == "__main__":
+    main()