Update quickstart.py for official release

quickstart.py

@@ -0,0 +1,129 @@
+import os
+import sys
+import torch
+from huggingface_hub import hf_hub_download
+from tokenizers import Tokenizer
+from safetensors.torch import load_file
+
+# Ensure UTF-8 console output on Windows
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding="utf-8")
+        sys.stdin.reconfigure(encoding="utf-8")
+    except Exception:
+        pass
+
+REPO_ID = "ViuAI/Sarus-500M"
+EOT_ID = 64002
+
+def main():
+    print("=" * 70)
+    print("🚀  ViuAI Sarus-500M -- 1-Click Interactive Terminal Chat")
+    print("=" * 70)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"[*] Compute Device: {device.upper()}")
+
+    # 1. Download or locate local architecture & tokenizer
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    
+    if os.path.exists(os.path.join(base_dir, "model.py")):
+        sys.path.insert(0, base_dir)
+    else:
+        print("[*] Downloading model architecture from Hugging Face...")
+        p_model = hf_hub_download(repo_id=REPO_ID, filename="model.py")
+        hf_hub_download(repo_id=REPO_ID, filename="config.py")
+        sys.path.insert(0, os.path.dirname(p_model))
+
+    from model import ViuAI
+    from config import ViuAIConfig
+
+    tok_local = os.path.join(base_dir, "tokenizer.json")
+    if not os.path.exists(tok_local):
+        print("[*] Downloading tokenizer...")
+        tok_local = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")
+
+    tokenizer = Tokenizer.from_file(tok_local)
+    cfg = ViuAIConfig(vocab_size=tokenizer.get_vocab_size(), context_length=2048)
+    model = ViuAI(cfg).to(device)
+
+    # 2. Download or locate Safetensors weights
+    weights_local = os.path.join(base_dir, "model.safetensors")
+    if not os.path.exists(weights_local):
+        print("[*] Downloading model.safetensors (~1.1 GB)...")
+        weights_local = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")
+
+    print(f"[*] Loading weights from {os.path.basename(weights_local)}...")
+    state_dict = load_file(weights_local)
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+
+    print("\n[OK] Sarus-500M Ready! Type 'exit' to quit.\n" + "-" * 70)
+
+    while True:
+        try:
+            prompt = input("\n👤 USER >> ").strip()
+            if not prompt:
+                continue
+            if prompt.lower() in ["exit", "quit", "q"]:
+                print("\nSession ended.")
+                break
+
+            formatted = f"<|user|>\n{prompt}<|endofturn|>\n<|assistant|>\n"
+            ids = tokenizer.encode(formatted).ids
+            idx = torch.tensor([ids], dtype=torch.long, device=device)
+
+            print("🤖 SARUS >> ", end="", flush=True)
+
+            generated = []
+            prev_text = ""
+
+            with torch.no_grad():
+                for _ in range(300):
+                    idx_cond = idx if idx.size(1) <= cfg.context_length else idx[:, -cfg.context_length:]
+                    logits, _ = model(idx_cond)
+                    logits = logits[:, -1, :]
+
+                    # Repetition penalty
+                    if generated:
+                        gen_t = torch.tensor(list(set(generated)), dtype=torch.long, device=device)
+                        score = logits[:, gen_t]
+                        logits[:, gen_t] = torch.where(score < 0, score * 1.15, score / 1.15)
+
+                    # Top-K & Top-P Sampling
+                    sampling_logits = logits / 0.15
+                    top_k_val = min(50, sampling_logits.size(-1))
+                    kth_vals, _ = torch.topk(sampling_logits, top_k_val)
+                    sampling_logits[sampling_logits < kth_vals[:, [-1]]] = float("-inf")
+
+                    sorted_logits, sorted_indices = torch.sort(sampling_logits, descending=True)
+                    cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+                    remove_mask = cum_probs > 0.90
+                    remove_mask[..., 1:] = remove_mask[..., :-1].clone()
+                    remove_mask[..., 0] = False
+                    indices_to_remove = torch.zeros_like(sampling_logits, dtype=torch.bool).scatter_(1, sorted_indices, remove_mask)
+                    sampling_logits[indices_to_remove] = float("-inf")
+
+                    probs = torch.softmax(sampling_logits, dim=-1)
+                    idx_next = torch.multinomial(probs, num_samples=1)
+                    token_id = idx_next.item()
+
+                    if token_id == EOT_ID:
+                        break
+
+                    generated.append(token_id)
+                    idx = torch.cat([idx, idx_next], dim=1)
+
+                    curr_text = tokenizer.decode(generated)
+                    chunk = curr_text[len(prev_text):]
+                    print(chunk, end="", flush=True)
+                    prev_text = curr_text
+
+            print()
+
+        except KeyboardInterrupt:
+            print("\nSession interrupted.")
+            break
+
+if __name__ == "__main__":
+    main()