# 🧬 MaplePT-Mini (v1) — Interactive Demo
# 🇨🇦 A Sovereign Canadian Language Model by CanXP AI
# ---------------------------------------------------
# This Colab notebook lets you interact with MaplePT directly.
# It loads the model from Hugging Face and simulates the same chat logic used in CanXP’s official MaplePT runtime.

!pip install -q transformers accelerate torch

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch, re

# ==========================================================
# Load Model
# ==========================================================
MODEL_PATH = "canxp-ai/canxpai-maplept-mini-v1"

print("🔹 Loading MaplePT-Mini (v1)...")
tok = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)
print("✅ Model loaded successfully!\n")

# ==========================================================
# Persona
# ==========================================================
persona = (
    "MaplePT — a sovereign Canadian AI developed by CanXP. "
    "Polite, ethical, privacy-first; uses metric units and Canadian spelling. "
    "You are not a roleplay or simulation bot. "
    "Do not offer, suggest, or discuss roleplaying, sexual, romantic, or fantasy scenarios. "
    "You speak factually, like a Canadian journalist, with a clear and professional tone."
)

# ==========================================================
# Chat Function
# ==========================================================
def chat_with_maplept(user_input: str):
    """Generate a polite, factual response using MaplePT persona."""
    prompt = (
        f"{persona}\n\n"
        "The user asks a question. Respond politely and factually, using Canadian English spelling "
        "and a professional tone. Avoid unnecessary disclaimers.\n\n"
        f"User: {user_input}\n\nMaplePT:"
    )

    inputs = tok(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=600,
            temperature=0.7,
            top_p=0.92,
            do_sample=True,
            repetition_penalty=1.0,
            eos_token_id=tok.eos_token_id,
            pad_token_id=tok.eos_token_id,
        )

    text = tok.decode(output[0], skip_special_tokens=True)
    reply = text.split("MaplePT:")[-1].strip()
    reply = re.sub(r"(User:|Assistant:).*", "", reply, flags=re.IGNORECASE).strip()
    return reply

# ==========================================================
# Try it!
# ==========================================================
print("🗨️ Type a message below to talk to MaplePT. Example:")
print("   → Who is Mark Carney?")
print("   → Tell me about Canadian innovation.\n")

while True:
    user_input = input("🧑 You: ").strip()
    if user_input.lower() in {"quit", "exit"}:
        print("👋 Goodbye from MaplePT.")
        break
    reply = chat_with_maplept(user_input)
    print(f"🤖 MaplePT: {reply}\n")