chat_jirack_27b.py
| 1 | #!/usr/bin/env python |
| 2 | # Chat для JiRackQwen38ForCausalLM (текущая версия архитектуры) |
| 3 | |
| 4 | import torch |
| 5 | from transformers import AutoTokenizer |
| 6 | from JiRackDeltaNet_27b import JiRackQwen38ForCausalLM |
| 7 | |
| 8 | MODEL_PATH = "model.pt" |
| 9 | TOKENIZER_DIR = "." |
| 10 | |
| 11 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 12 | print(f"Loading on {device}...") |
| 13 | |
| 14 | model = JiRackQwen38ForCausalLM.load_checkpoint(MODEL_PATH, dtype=torch.bfloat16) |
| 15 | model.to(device).eval() |
| 16 | |
| 17 | tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR, trust_remote_code=True) |
| 18 | |
| 19 | print("Ready! Type your message:") |
| 20 | history = [] |
| 21 | |
| 22 | while True: |
| 23 | user_input = input("\nYou: ").strip() |
| 24 | if not user_input: |
| 25 | continue |
| 26 | if user_input.lower() in ["exit", "quit"]: |
| 27 | break |
| 28 | |
| 29 | history.append({"role": "user", "content": user_input}) |
| 30 | |
| 31 | input_ids = tokenizer.apply_chat_template( |
| 32 | history, |
| 33 | add_generation_prompt=True, |
| 34 | return_tensors="pt" |
| 35 | ) |
| 36 | if not torch.is_tensor(input_ids): |
| 37 | input_ids = input_ids["input_ids"] |
| 38 | |
| 39 | input_ids = input_ids.to(device) |
| 40 | |
| 41 | with torch.no_grad(): |
| 42 | # Generate token-by-token |
| 43 | max_new = 256 |
| 44 | for _ in range(max_new): |
| 45 | logits = model(input_ids) |
| 46 | next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) |
| 47 | input_ids = torch.cat([input_ids, next_token], dim=1) |
| 48 | |
| 49 | if next_token.item() == tokenizer.eos_token_id: |
| 50 | break |
| 51 | |
| 52 | reply = tokenizer.decode(input_ids[0, input_ids.shape[1]-max_new:], skip_special_tokens=True) |
| 53 | print(f"\nAssistant: {reply}") |
| 54 | history.append({"role": "assistant", "content": reply}) |
| 55 | |