Update chat.py

chat.py

@@ -2,11 +2,12 @@ import torch
 import tiktoken
 from model import GPTConfig, GPT
 
+# --- Config ---
 ckpt_path = '/home/user/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)
@@ -23,7 +24,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: ")
@@ -37,10 +38,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__":