Update README.md

README.md

@@ -129,17 +129,57 @@ This reinforces the distinction between:
 
 ```python
 import torch
-from safetensors.torch import load_file
-from model import NanoThink
 from tokenizers import Tokenizer
+from model import NanoThink
+from safetensors.torch import load_file
+
+MODEL_PATH = "model.safetensors"
+TOKENIZER_PATH = "tokenizer.json"
 
-tokenizer = Tokenizer.from_file("tokenizer.json")
 
-model = NanoThink(vocab_size=1229)
-state_dict = load_file("model.safetensors")
-model.load_state_dict(state_dict)
+tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
 
+model = NanoThink(vocab_size=tokenizer.get_vocab_size())
+model.load_state_dict(load_file(MODEL_PATH))
 model.eval()
+
+history = ""
+
+while True:
+    user_input = input("You: ")
+
+    if user_input.lower() in ["get out", "exit", "quit"]:
+        break
+
+    prompt = history + f"\n<USER>\n{user_input}\n</USER>\n"
+
+    input_ids = torch.tensor([tokenizer.encode(prompt).ids])
+
+    output_ids = []
+
+    for _ in range(120):
+        logits = model(input_ids)
+        next_token = torch.multinomial(torch.softmax(logits[0, -1], dim=-1), 1).item()
+
+        input_ids = torch.cat([input_ids, torch.tensor([[next_token]])], dim=1)
+        output_ids.append(next_token)
+
+        text = tokenizer.decode(output_ids)
+
+        if "</ANSWER>" in text:
+            break
+
+    output = tokenizer.decode(output_ids)
+
+    
+    if "<ANSWER>" in output:
+        output = output.split("<ANSWER>")[1].split("</ANSWER>")[0]
+
+    print("\n💬 Answer:")
+    print(output.strip())
+    print("\n" + "-"*50 + "\n")
+
+    history += f"\n<USER>\n{user_input}\n</USER>\n<ANSWER>\n{output.strip()}\n</ANSWER>\n"
 ```
 
 ---