# ==============================================================================
# JiRack 7B Chat (JiRackPrecision tokenizer)  HF CMSManhattan/JiRackPrecisionTokenizer
# COPYRIGHT (c) 2026 Konstantin Vladimirovich Grabko.
#
# Adapted from chat_10b_v6_st_pt.py for the JiRack 7B checkpoint:
#   * model class:    JiRackTernaryUltra_7b.JiRackTransformer / JiRackConfig
#   * checkpoint fmt: dict with a "model" key (not "model_state_dict"),
#                     produced by convert_ds7b_to_jirack.py
#   * tokenizer:      YOUR extended DeepSeek-Qwen tokenizer folder (adds the
#                     110 custom tags on top of the original DeepSeek vocab),
#                     NOT the Llama-3 tokenizer used by the old 10B model.
#   * uses tokenizer.apply_chat_template() (chat_template.jinja is already
#     sitting in that folder) instead of a hand-rolled Llama-3 prompt string.
#
# CMS Manhattan JiRack Technology — PATENT PENDING
#
# This code is proprietary.
# Personal and non-commercial research use is allowed.
# Any commercial use, derivative works for profit, or distribution
# requires a paid license and 5% royalty.
#
# Unauthorized commercial use is strictly prohibited.
# Contact: grabko@cmsmanhattan.com
# =============================================================================


import os
import sys
import torch
from transformers import AutoTokenizer

sys.path.append(os.getcwd())
from JiRackTernaryUltra_7b import JiRackTransformer, JiRackConfig

# ========================= EDIT THESE =========================
MODEL_PATH    = "model.pt"
TOKENIZER_DIR = "."   # your extended tokenizer folder
NO_THINK      = True   # True = skip <think> reasoning, answer directly
# ================================================================


def load_model(model_path: str):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"🚀 Загрузка модели на устройство: {device.upper()}")

    config = JiRackConfig()
    model = JiRackTransformer(config, use_checkpoint=False)

    print(f"📥 Загрузка весов из {model_path}...")
    try:
        ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
        state_dict = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt

        missing, unexpected = model.load_state_dict(state_dict, strict=False)
        # lambda_ buffers legitimately absent from some older checkpoints
        real_missing = [k for k in missing if not k.endswith("lambda_")]
        if real_missing:
            print(f"⚠️ Пропущено ключей: {len(real_missing)} -> {real_missing[:10]}")
        if unexpected:
            print(f"⚠️ Лишние ключи: {len(unexpected)} -> {unexpected[:10]}")
    except Exception as e:
        print(f"❌ Критическая ошибка при загрузке весов: {e}")
        sys.exit(1)

    model = model.to(dtype=torch.bfloat16, device=device).eval()
    model.set_lambda(0.0)  # full-precision fast path, no fake-quant at inference

    if device == "cuda":
        vram = torch.cuda.memory_allocated(0) / 1024**3
        print(f"✅ VRAM занято: {vram:.1f} GB")
    else:
        print("⚠️ ВНИМАНИЕ: Запуск на CPU будет очень медленным.")

    print("✅ Модель успешно загружена.")
    return model, device


@torch.no_grad()
def generate_text(model, tokenizer, input_ids, stop_tokens, max_new_tokens=512, device="cuda"):
    curr_ids = input_ids.to(device)
    prompt_len = curr_ids.shape[1]
    printed = ""

    temperature = 0.6          # DeepSeek-R1 distill models recommend ~0.5-0.7
    top_p = 0.95
    repetition_penalty = 1.15

    print("JiRack: ", end="", flush=True)

    for _ in range(max_new_tokens):
        with torch.autocast(device_type=("cuda" if device == "cuda" else "cpu"), dtype=torch.bfloat16):
            logits = model(curr_ids)
            next_token_logits = logits[:, -1, :].float() / temperature

        # Repetition penalty
        for token_id in set(curr_ids[0].tolist()):
            if next_token_logits[0, token_id] < 0:
                next_token_logits[0, token_id] *= repetition_penalty
            else:
                next_token_logits[0, token_id] /= repetition_penalty

        # Top-p sampling
        sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        next_token_logits[0, sorted_indices[sorted_indices_to_remove]] = -float('Inf')
        probs = torch.softmax(next_token_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)

        curr_ids = torch.cat([curr_ids, next_token], dim=1)
        # decode the whole generated tail each step and print only the new part;
        # keeps multi-token UTF-8 chars (emoji etc.) intact instead of \ufffd
        decoded = tokenizer.decode(curr_ids[0, prompt_len:], skip_special_tokens=True)
        if not decoded.endswith("\ufffd"):
            print(decoded[len(printed):], end="", flush=True)
            printed = decoded

        if next_token.item() in stop_tokens:
            break

    print("\n")
    return curr_ids


def main():
    if not os.path.exists(MODEL_PATH):
        print(f"❌ Файл {MODEL_PATH} не найден!")
        return

    try:
        tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
    except Exception as e:
        print(f"❌ Ошибка токенайзера: {e}")
        return

    model, device = load_model(MODEL_PATH)

    # Stop tokens: eos + any DeepSeek "end" style special tokens present.
    stop_tokens = set()
    if tokenizer.eos_token_id is not None:
        stop_tokens.add(tokenizer.eos_token_id)
    for name in ("<|end_of_sentence|>", "<|endoftext|>", "<|im_end|>"):
        tid = tokenizer.convert_tokens_to_ids(name)
        if tid is not None and tid != tokenizer.unk_token_id:
            stop_tokens.add(tid)

    print("\n" + "=" * 80)
    print("✅ JiRack 7B (DeepSeek-R1-Distill-Qwen, extended tokenizer) Ready")
    print("=" * 80 + "\n")

    history = []  # list of {"role": ..., "content": ...} for multi-turn chat

    while True:
        try:
            user_input = input("User: ")
            if user_input.lower() in ["exit", "quit", "q"]:
                break
            if not user_input.strip():
                continue

            history.append({"role": "user", "content": user_input})

            # Uses chat_template.jinja already present in TOKENIZER_DIR.
            input_ids = tokenizer.apply_chat_template(
                history,
                add_generation_prompt=True,
                return_tensors="pt",
                return_dict=False,
            )
            # some transformers versions return a BatchEncoding here regardless;
            # unwrap it defensively so we always end up with a plain tensor
            if not torch.is_tensor(input_ids):
                input_ids = input_ids["input_ids"]

            # NO_THINK: chat_template ends with '<|Assistant|><think>\n'.
            # Appending '</think>\n\n' makes the model skip reasoning and
            # answer directly (standard trick for R1-distill models).
            if NO_THINK:
                close_ids = tokenizer.encode("</think>\n\n", add_special_tokens=False, return_tensors="pt")
                input_ids = torch.cat([input_ids, close_ids], dim=1)

            curr_ids = generate_text(model, tokenizer, input_ids, stop_tokens, device=device)

            new_tokens = curr_ids[0, input_ids.shape[1]:]
            reply = tokenizer.decode(new_tokens, skip_special_tokens=True)
            history.append({"role": "assistant", "content": reply})

        except KeyboardInterrupt:
            print("\nStopped.")
            break
        except Exception:
            import traceback
            print("\n❌ Ошибка:")
            traceback.print_exc()


if __name__ == "__main__":
    main()