Update README.md

README.md

@@ -13,10 +13,13 @@ metrics:
 
 Welcome to SmaLLMPro 350M, our latest Instruct-Model based on FineWeb-Edu.
 
-# 1. Parameters
-SmaLLMPro 350M has 353,550,000 parameters.
-1. Decayed parameter tensors: 98, with 354,549,760 parameters
-2. Non-decayed parameter tensors: 49, with 50,176 parameters
+# 1. Model Details
+- **Parameters:** 353.55M
+- **Layers:** 24
+- **Heads:** 16
+- **Embedding Dim:** 1024
+- **Context Length:** 1024
+- **Format:** ONNX (Opset 18)
 
 # 2. Trainingcode
 ```python
@@ -707,11 +710,12 @@ import torch
 import tiktoken
 from model import GPTConfig, GPT
 
-ckpt_path = '/home/user/350m_SmaLLMPro_Final/SmaLLMPro_iter_1500.pt'
+# --- Config ---
+ckpt_path = '/media/leo/Data/checkpoints/350m_SmaLLMPro_Final/SmaLLMPro_iter_1500.pt'
 device = 'cuda'
 enc = tiktoken.get_encoding("gpt2")
 
-print("Loading SmaLLMPro model checkpoint...")
+print("Loading SmaLLMPro...")
 checkpoint = torch.load(ckpt_path, map_location=device)
 gptconf = GPTConfig(**checkpoint['model_args'])
 model = GPT(gptconf)
@@ -728,7 +732,7 @@ model.to(device)
 print(f"Model {ckpt_path} ready!\n")
 
 def run_chat():
-    print("--- SmaLLMPro Chatbot (Type 'exit' to leave) ---")
+    print("--- SmaLLMPro Chatbot (Type 'exit' to quit) ---")
     
     while True:
         user_input = input("You: ")
@@ -742,10 +746,16 @@ def run_chat():
         print("SmaLLMPro: ", end="", flush=True)
         with torch.no_grad():
             with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
-                y = model.generate(x, max_New_tokens=256, temperature=0.4, top_k=25)
+                y = model.generate(x, max_new_tokens=500, temperature=0.65, top_k=25)
+                
                 full_text = enc.decode(y[0].tolist())
                 
-                response = full_text.split("Response:\n")[-1].split("<|endoftext|>")[0].strip()
+                if "Response:\n" in full_text:
+                    response = full_text.split("Response:\n")[-1]
+                else:
+                    response = full_text
+                
+                response = response.split("<|endoftext|>")[0].split("Instruction:")[0].strip()
                 print(response + "\n")
 
 if __name__ == "__main__":