Update chat.py

chat.py

@@ -1,3 +1,8 @@
+"""
+Chat interface for CosmicFish model downloaded from Hugging Face Hub.
+Uses safetensors format only for secure model loading.
+"""
+
 import os
 import sys
 import time
@@ -24,6 +29,15 @@ except ImportError:
     print("Install with: pip install transformers huggingface-hub")
     sys.exit(1)
 
+# Required for safetensors
+try:
+    from safetensors.torch import load_file
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+    print("Safetensors not available. Install with: pip install safetensors")
+    sys.exit(1)
+
 # Set up logging
 logging.basicConfig(
     level=logging.INFO,
@@ -392,24 +406,16 @@ class CosmicFishChatSession:
             "\nUser:"
         ]
 
-        # Print welcome message
         if config.display_welcome:
             self._print_welcome_message()
 
     def _print_welcome_message(self):
-        """Print a welcome message to the user."""
         welcome_text = f"""
 {'=' * 80}
-Welcome to CosmicFish!
-
-This is a {self.model.get_num_params() / 1e6:.1f}M parameter model made by MistyozAI.
-CosmicFish features advanced architecture with RoPE, GQA, SwiGLU, and RMSNorm.
-
-⚠️ DISCLAIMER: Since this {self.model.get_num_params() / 1e6:.1f}M parameter model is relatively
-small, it is more likely to give incorrect answers or hallucinate compared to
-larger models. Please verify important information from reliable sources.
+Welcome to CosmicFish chat interface
 
-Model: {DEFAULT_MODEL_REPO}
+This is a {self.model.get_num_params() / 1e6:.1f}M parameter model.
+CosmicFish is an efficient LLM with an advanced architecture.
 
 Type your prompts and CosmicFish will respond.
 
@@ -423,6 +429,14 @@ Special commands:
 - /temp [value]: Set temperature (between 0.1 and 2.0)
 - /penalty [value]: Set repetition penalty (1.0-2.0)
 - /debug: Toggle debug mode
+
+
+Note: CosmicFIsh may generate incorrect or fictional responses. Verify facts if needed.
+
+Visit https://cosmicfish.ai for more info
+
+
+Developed by Mistyoz AI (https://www.mistyoz.com)
 {'=' * 80}
 """
         print(colored(welcome_text, 'cyan'))
@@ -494,19 +508,15 @@ Special commands:
         return False
 
     def _clean_token_text(self, text):
-
         text = text.replace('��', "'")
-
         text = text.replace('�', "'")
         text = text.replace('\ufffd', "'")
         text = text.replace('\uFFFD', "'")
-
-        text = text.replace('’', "'")
-        text = text.replace('“', "'")
-        text = text.replace('�', "'")
-        text = text.replace('â€"', "'")
-        text = text.replace('â€"', "'")
-
+        text = text.replace('’', "'")
+        text = text.replace('â€Å"', "'")
+        text = text.replace('�', "'")
+        text = text.replace('â€"', "'")
+        text = text.replace('â€"', "'")
         return text
 
     def generate_with_repetition_penalty(self, input_ids, max_new_tokens, temperature, top_k, penalty=1.2, live=False):
@@ -772,6 +782,7 @@ Token usage statistics:
 - Current temperature: {self.config.temperature}
 - Model: CosmicFish ({self.model.get_num_params() / 1e6:.1f}M parameters)
 - Source: {DEFAULT_MODEL_REPO}
+- Format: Safetensors (secure)
 """
             print(colored(stats, 'yellow'))
             return True
@@ -910,7 +921,7 @@ Token usage statistics:
 
 
 def download_cosmicfish_from_hub(model_repo=DEFAULT_MODEL_REPO, device='cpu'):
-    """Download and load CosmicFish model from Hugging Face Hub"""
+    """Download and load CosmicFish model from Hugging Face Hub (safetensors only)"""
     print(colored(f"Downloading CosmicFish from Hugging Face: {model_repo}", "cyan"))
 
     try:
@@ -945,10 +956,19 @@ def download_cosmicfish_from_hub(model_repo=DEFAULT_MODEL_REPO, device='cpu'):
         print("Creating model...")
         model = CosmicFish(config)
 
-        # Load weights
-        print("Loading weights...")
-        weights_path = os.path.join(cache_dir, "pytorch_model.bin")
-        state_dict = torch.load(weights_path, map_location=device)
+        # Load weights from safetensors ONLY
+        print("Loading weights from safetensors...")
+        safetensors_path = os.path.join(cache_dir, "model.safetensors")
+
+        if not os.path.exists(safetensors_path):
+            raise FileNotFoundError(f"model.safetensors not found in {cache_dir}. This model requires safetensors format.")
+
+        state_dict = load_file(safetensors_path)
+
+        # Handle weight sharing: lm_head.weight shares with transformer.wte.weight
+        if 'lm_head.weight' not in state_dict and 'transformer.wte.weight' in state_dict:
+            state_dict['lm_head.weight'] = state_dict['transformer.wte.weight']
+
         model.load_state_dict(state_dict)
         model.to(device)
         model.eval()
@@ -964,14 +984,12 @@ def download_cosmicfish_from_hub(model_repo=DEFAULT_MODEL_REPO, device='cpu'):
 
 
 def load_tokenizer():
-    print("Loading tokenizer...")
     tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    print("Tokenizer loaded")
     return tokenizer
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Chat with CosmicFish model from Hugging Face Hub")
+    parser = argparse.ArgumentParser(description="Chat with CosmicFish")
 
     # Model parameters
     parser.add_argument("--model_repo", type=str, default=DEFAULT_MODEL_REPO,
@@ -982,7 +1000,7 @@ def main():
     # Generation parameters
     parser.add_argument("--temperature", type=float, default=0.7,
                         help="Temperature for sampling (default: 0.7)")
-    parser.add_argument("--max_tokens", type=int, default=512,
+    parser.add_argument("--max_tokens", type=int, default=1024,
                         help="Maximum number of tokens to generate per response")
     parser.add_argument("--min_tokens", type=int, default=10,
                         help="Minimum number of tokens to generate per response")
@@ -1049,7 +1067,7 @@ def main():
         chat = CosmicFishChatSession(model, tokenizer, config)
 
         # Main chat loop
-        print(colored("\nCosmicFish initialized! Type your message (or /help for commands).\n", 'cyan'))
+        print(colored("\nCosmicFish initialized from Hugging Face! Type your message (or /help for commands).\n", 'cyan'))
 
         while True:
             try:
@@ -1087,8 +1105,6 @@ def main():
                             if not live_buffer:
                                 print(final_response, end="")
                             break
-
-                        # If we have a token to display
                         if token:
                             # Check if token contains <|endoftext|> and remove it if present
                             if "<|endoftext|>" in token: