updated chat.py for better inference

chat.py

@@ -11,7 +11,7 @@ import torch
 import numpy as np
 from termcolor import colored
 import logging
-import readline  # Enables arrow key history in terminal input
+import readline
 import re
 import textwrap
 import random
@@ -25,7 +25,7 @@ try:
     HF_AVAILABLE = True
 except ImportError:
     HF_AVAILABLE = False
-    print("❌ Required libraries not available.")
+    print("Required libraries not available.")
     print("Install with: pip install transformers huggingface-hub")
     sys.exit(1)
 
@@ -38,7 +38,7 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # Default model repository
-DEFAULT_MODEL_REPO = "Mistyoz-AI/CosmicFish-120M"
+DEFAULT_MODEL_REPO = "MistyozAI/CosmicFish-120M"
 
 # Default prompt template
 DEFAULT_PROMPT_TEMPLATE = "Below is a conversation between a helpful AI assistant and a human. The assistant is knowledgeable, friendly, and provides detailed and accurate responses.\n\n"
@@ -54,7 +54,7 @@ class CosmicConfig:
                  n_head=16,
                  n_embd=704,
                  bias=True,
-                 dropout=0.0,  # Always 0 for inference
+                 dropout=0.0,
                  n_query_groups=4,
                  eps=1e-6,
                  use_rotary=True,
@@ -405,11 +405,15 @@ class CosmicFishChatSession:
         """Print a welcome message to the user."""
         welcome_text = f"""
 {'=' * 80}
-Welcome to CosmicFish chat interface (Hugging Face Hub)
+Welcome to CosmicFish!
 
-This is a {self.model.get_num_params() / 1e6:.1f}M parameter model loaded from HF Hub.
+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.
+
 Model: {DEFAULT_MODEL_REPO}
 
 Type your prompts and CosmicFish will respond.
@@ -495,9 +499,19 @@ Special commands:
         return False
 
     def _clean_token_text(self, text):
-        """Clean token text by fixing encoding issues."""
-        # Fix the specific issue with �� -> '
+
         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('â€"', "'")
+
         return text
 
     def generate_with_repetition_penalty(self, input_ids, max_new_tokens, temperature, top_k, penalty=1.2, live=False):
@@ -902,13 +916,13 @@ Token usage statistics:
 
 def download_cosmicfish_from_hub(model_repo=DEFAULT_MODEL_REPO, device='cpu'):
     """Download and load CosmicFish model from Hugging Face Hub"""
-    print(colored(f"🤗 Downloading CosmicFish from Hugging Face Hub: {model_repo}", "cyan"))
+    print(colored(f"Downloading CosmicFish from Hugging Face: {model_repo}", "cyan"))
 
     try:
         # Download the model files to local cache
-        print("📥 Downloading model files...")
+        print("Downloading model files...")
         cache_dir = snapshot_download(repo_id=model_repo, cache_dir=None)
-        print(f"✅ Model cached at: {cache_dir}")
+        print(f"Model cached at: {cache_dir}")
 
         # Load config
         config_path = os.path.join(cache_dir, "config.json")
@@ -933,32 +947,31 @@ def download_cosmicfish_from_hub(model_repo=DEFAULT_MODEL_REPO, device='cpu'):
         )
 
         # Create model
-        print("🧠 Creating model...")
+        print("Creating model...")
         model = CosmicFish(config)
 
         # Load weights
-        print("⚖️ Loading weights...")
+        print("Loading weights...")
         weights_path = os.path.join(cache_dir, "pytorch_model.bin")
         state_dict = torch.load(weights_path, map_location=device)
         model.load_state_dict(state_dict)
         model.to(device)
         model.eval()
 
-        print(f"✅ Model loaded: {model.get_num_params() / 1e6:.1f}M parameters")
-        print(f"🎯 Device: {device}")
+        print(f"Model loaded: {model.get_num_params() / 1e6:.1f}M parameters")
+        print(f"Device: {device}")
         return model, config
 
     except Exception as e:
-        print(colored(f"❌ Error downloading/loading model: {str(e)}", "red"))
-        print(colored("💡 Make sure you have internet connection and the model repo exists", "yellow"))
+        print(colored(f"Error downloading/loading model: {str(e)}", "red"))
+        print(colored("Make sure you have internet connection and the model repo exists", "yellow"))
         sys.exit(1)
 
 
 def load_tokenizer():
-    """Load GPT-2 tokenizer"""
-    print("🔤 Loading GPT-2 tokenizer...")
+    print("Loading tokenizer...")
     tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    print("✅ Tokenizer loaded")
+    print("Tokenizer loaded")
     return tokenizer
 
 
@@ -1007,7 +1020,7 @@ def main():
     # Configure device
     device = args.device
     if device == "cuda" and not torch.cuda.is_available():
-        print(colored("⚠️ CUDA is not available, falling back to CPU", "yellow"))
+        print(colored("CUDA is not available, falling back to CPU", "yellow"))
         device = "cpu"
 
     try:
@@ -1041,7 +1054,7 @@ def main():
         chat = CosmicFishChatSession(model, tokenizer, config)
 
         # Main chat loop
-        print(colored("\n🚀 CosmicFish initialized from Hugging Face Hub. Type your message (or /help for commands).\n", 'cyan'))
+        print(colored("\nCosmicFish initialized! Type your message (or /help for commands).\n", 'cyan'))
 
         while True:
             try: