Update README.md

README.md

@@ -3,22 +3,43 @@ license: apache-2.0
 tags:
 - text-generation
 - language-model
-- LLM
-- CosmicFish
-- 120M
+- causal-lm
+- cosmicfish
+- 120m
 - transformer
+- rope
+- gqa
+- swiglu
+- rmsnorm
 language: en
 datasets:
 - CosmicSet-1.0
-- akkiisfrommars/TreeCorpusCleaned
+- akkiisfrommars/TreeCorpusCleanedmodel
 model_type: CosmicFish
+pipeline_tag: text-generation
 ---
 
-
 # CosmicFish-120M
 
 A 120M parameter language model with modern architecture improvements developed by Mistyoz AI.
 
+## Quick Start
+
+**The easiest way to chat with CosmicFish is using our chat.py script:**
+
+```bash
+# Download the chat script from this repository
+wget https://huggingface.co/MistyozAI/CosmicFish-120M/resolve/main/chat.py
+
+# Install dependencies
+pip install transformers huggingface-hub termcolor
+
+# Run the chat interface (automatically downloads model)
+python chat.py
+```
+
+The `chat.py` script handles all model loading, generation, and provides the best chat experience with live streaming, repetition penalty, and conversation commands.
+
 ## Model Details
 
 - **Parameters**: 121M
@@ -27,6 +48,7 @@ A 120M parameter language model with modern architecture improvements developed
 - **Vocabulary**: 50,257 tokens
 - **Training Data**: CosmicSet 1.0
 - **Developer**: Mistyoz AI
+- **Repository**: MistyozAI/CosmicFish-120M
 
 ## Usage
 
@@ -36,52 +58,62 @@ A 120M parameter language model with modern architecture improvements developed
 pip install transformers huggingface-hub termcolor
 ```
 
-### Loading the Model
+### Quick Chat Interface
 
 ```python
+from transformers import GPT2Tokenizer
+from huggingface_hub import snapshot_download
 import torch
 import json
-from transformers import GPT2Tokenizer
-from modeling_cosmicfish import CosmicFish, CosmicConfig
-
-# Load model
-with open("config.json") as f:
-    config_dict = json.load(f)
+import os
 
-config = CosmicConfig(**{k: v for k, v in config_dict.items() if k in [
-    'vocab_size', 'block_size', 'n_layer', 'n_head', 'n_embd', 'bias',
-    'use_rotary', 'use_swiglu', 'use_gqa', 'n_query_groups'
-]})
-config.dropout = 0.0  # Inference mode
-
-model = CosmicFish(config)
-model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
-model.eval()
+# Download model from Hugging Face Hub
+cache_dir = snapshot_download(repo_id="MistyozAI/CosmicFish-120M")
 
 # Load tokenizer
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+# Load config
+with open(os.path.join(cache_dir, "config.json")) as f:
+    config_dict = json.load(f)
+
+# Load model weights
+state_dict = torch.load(os.path.join(cache_dir, "pytorch_model.bin"), map_location="cpu")
+
+# Note: Full model class available in the repository
+print("Model downloaded and ready for use!")
 ```
 
-### Basic Generation
+### Advanced Generation with Repetition Penalty
 
 ```python
-def generate_text(prompt, max_tokens=100):
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
+def generate_with_repetition_penalty(model, tokenizer, prompt, max_tokens=100, temperature=0.7, penalty=1.2):
+    input_ids = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
+    generated = input_ids.clone()
     
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs,
-            max_new_tokens=max_tokens,
-            temperature=0.7,
-            top_k=40,
-            do_sample=True
-        )
+    for _ in range(max_tokens):
+        with torch.no_grad():
+            logits, _ = model(generated)
+        
+        next_token_logits = logits[:, -1, :] / temperature
+        
+        # Apply repetition penalty
+        if penalty > 1.0:
+            for token_id in set(generated[0].tolist()):
+                if next_token_logits[0, token_id] > 0:
+                    next_token_logits[0, token_id] /= penalty
+                else:
+                    next_token_logits[0, token_id] *= penalty
+        
+        probs = torch.nn.functional.softmax(next_token_logits, dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1)
+        
+        if next_token.item() == tokenizer.eos_token_id:
+            break
+            
+        generated = torch.cat([generated, next_token], dim=1)
     
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Example
-text = generate_text("The future of AI is")
-print(text)
+    return tokenizer.decode(generated[0], skip_special_tokens=True)
 ```
 
 ### Chat Interface
@@ -95,31 +127,21 @@ def chat_with_model():
         if user_input.lower() in ['quit', 'exit']:
             break
         
-       
         context = "Below is a conversation between a human and an AI assistant.\n\n"
         for human, ai in conversation:
             context += f"Human: {human}\nAssistant: {ai}\n\n"
         context += f"Human: {user_input}\nAssistant:"
         
-        # Generate response
-        inputs = tokenizer.encode(context, return_tensors="pt")
-        if inputs.shape[1] > 400: 
-            inputs = inputs[:, -400:]
-        
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs,
-                max_new_tokens=150,
-                temperature=0.7,
-                top_k=40,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
-            )
-        
-        response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-        response = response.split('\n')[0].strip()
+        # Generate response with repetition penalty
+        response = generate_with_repetition_penalty(
+            model, tokenizer, context, 
+            max_tokens=150, temperature=0.7, penalty=1.2
+        )
         
+        # Extract just the assistant's response
+        response = response.split("Assistant:")[-1].split('\n')[0].strip()
         print(f"CosmicFish: {response}")
+        
         conversation.append((user_input, response))
 
 chat_with_model()
@@ -130,7 +152,7 @@ chat_with_model()
 CosmicFish uses several modern improvements over standard transformers:
 
 - **RoPE (Rotary Position Embeddings)**: Better position encoding than absolute positions
-- **GQA (Grouped-Query Attention)**: Reduces memory usage with 4 query groups  
+- **GQA (Grouped-Query Attention)**: Reduces memory usage with 4 query groups 
 - **SwiGLU**: More effective activation function than ReLU/GELU
 - **RMSNorm**: Simpler, more stable normalization than LayerNorm
 
@@ -148,6 +170,7 @@ CosmicFish uses several modern improvements over standard transformers:
 - **File Size**: 243MB
 
 ## Limitations
+
 - Small model size (120M parameters) may produce less accurate responses
 - 512 token context limit
 - Training data cutoff applies
@@ -158,6 +181,6 @@ CosmicFish uses several modern improvements over standard transformers:
 
 Apache 2.0 - see LICENSE file.
 
-
 ## Credit
+
 If you use CosmicFish-120M, please credit Mistyoz AI.