| import os |
| import sys |
| import torch |
| from huggingface_hub import hf_hub_download |
| from tokenizers import Tokenizer |
|
|
| |
| if sys.platform == "win32": |
| try: |
| sys.stdout.reconfigure(encoding="utf-8") |
| sys.stdin.reconfigure(encoding="utf-8") |
| except Exception: |
| pass |
|
|
| REPO_ID = "ViuAI/ViuAI-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()}") |
|
|
| |
| base_dir = os.path.dirname(os.path.abspath(__file__)) |
| code_dir = os.path.join(base_dir, "code") |
| |
| if os.path.exists(os.path.join(code_dir, "model.py")): |
| sys.path.insert(0, code_dir) |
| else: |
| print("[*] Downloading model architecture from Hugging Face...") |
| p_model = hf_hub_download(repo_id=REPO_ID, filename="code/model.py") |
| p_cfg = hf_hub_download(repo_id=REPO_ID, filename="code/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", "tokenizer.json") |
| if not os.path.exists(tok_local): |
| 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) |
|
|
| |
| ckpt_local = os.path.join(base_dir, "dpo_checkpoints", "dpo_5k", "dpo_5k_final.pt") |
| if not os.path.exists(ckpt_local): |
| print("[*] Downloading DPO-5K Final Weights (~944 MB)...") |
| ckpt_local = hf_hub_download(repo_id=REPO_ID, filename="dpo_checkpoints/dpo_5k/dpo_5k_final.pt") |
|
|
| print(f"[*] Loading weights from {os.path.basename(ckpt_local)}...") |
| state = torch.load(ckpt_local, map_location=device, weights_only=True) |
| raw_sd = state.get("model", state.get("model_state_dict", state)) |
| clean_sd = {k.replace("_orig_mod.", "").replace("module.", ""): v for k, v in raw_sd.items()} |
| model.load_state_dict(clean_sd, strict=False) |
| model.eval() |
|
|
| print("\n[OK] Model 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, :] |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
