app.py

import os
import gradio as gr
import torch
import json
import re
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
from concurrent.futures import ThreadPoolExecutor
import logging
import time
import threading

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_ID = "YoussefElsafi/Aiko-350M"

# ── CPU threading ─────────────────────────────────────────────────────────────
PHYSICAL_CORES   = os.cpu_count() or 2
NUM_MODEL_COPIES = 2
MAX_CONCURRENT   = NUM_MODEL_COPIES
THREADS_PER_GEN  = max(1, PHYSICAL_CORES // NUM_MODEL_COPIES)

torch.set_num_threads(THREADS_PER_GEN)
torch.set_num_interop_threads(2)

for var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS",
            "VECLIB_MAXIMUM_THREADS", "NUMEXPR_NUM_THREADS"):
    os.environ[var] = str(THREADS_PER_GEN)

logger.info(f"CPU cores: {PHYSICAL_CORES} | Model copies: {NUM_MODEL_COPIES} | "
            f"Concurrent: {MAX_CONCURRENT} | Threads/gen: {THREADS_PER_GEN}")

# ── Model loading ─────────────────────────────────────────────────────────────
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

if torch.cuda.is_available():
    DEVICE = "cuda"
    DTYPE  = torch.float16
    logger.info("Device: CUDA — float16")
else:
    DEVICE = "cpu"
    DTYPE  = torch.float32
    logger.info("Device: CPU — float32 (INT8 quantization will apply)")

def load_one_model(idx: int):
    logger.info(f"Loading model copy #{idx}...")
    m = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        dtype             = DTYPE,
        device_map        = "auto" if DEVICE == "cuda" else "cpu",
        trust_remote_code = True,
        low_cpu_mem_usage = True,
    )
    m.eval()

    if DEVICE == "cpu":
        try:
            m = torch.quantization.quantize_dynamic(
                m, {torch.nn.Linear}, dtype=torch.qint8,
            )
            logger.info(f"Model #{idx}: INT8 quantization applied ✓")
        except Exception as e:
            logger.warning(f"Model #{idx}: quantization skipped: {e}")

    if hasattr(torch, "compile"):
        try:
            m = torch.compile(m, mode="default")
            logger.info(f"Model #{idx}: torch.compile applied ✓")
        except Exception as e:
            logger.warning(f"Model #{idx}: torch.compile skipped: {e}")

    return m

print(f"Loading {NUM_MODEL_COPIES} model copies...")
models = [load_one_model(i) for i in range(NUM_MODEL_COPIES)]
model_locks = [threading.Lock() for _ in range(NUM_MODEL_COPIES)]
print(f"All {NUM_MODEL_COPIES} model copies loaded!")

# ── Aiko config ───────────────────────────────────────────────────────────────
JSON_PREFIX = '{"internal_monologue":"'

GEN_CONFIG = dict(
    max_new_tokens     = 400,
    do_sample          = True,
    temperature        = 0.85,
    top_p              = 0.9,
    repetition_penalty = 1.05,
    pad_token_id       = tokenizer.pad_token_id,
    eos_token_id       = tokenizer.eos_token_id,
    use_cache          = True,
)

# ── Concurrency ───────────────────────────────────────────────────────────────
executor      = ThreadPoolExecutor(max_workers=MAX_CONCURRENT, thread_name_prefix="aiko")
queue_counter = threading.Semaphore(8)

def acquire_free_model(timeout: float = 60.0):
    deadline = time.time() + timeout
    while time.time() < deadline:
        for i, lock in enumerate(model_locks):
            if lock.acquire(blocking=False):
                return i, models[i]
        time.sleep(0.05)
    return None, None

# ── Auth ──────────────────────────────────────────────────────────────────────
ALLOWED_TOKEN = os.environ.get("AIKO_API_KEYS", "")

def validate_key(api_key: str) -> bool:
    if not ALLOWED_TOKEN:
        return True
    return api_key.strip() == ALLOWED_TOKEN.strip()

# ── JSON extraction ───────────────────────────────────────────────────────────
def extract_aiko_data(text):
    try:
        return json.loads(text)
    except:
        pass
    
    try:
        end_idx = text.rindex('}') + 1
        return json.loads(text[:end_idx])
    except:
        pass
    
    data = {}
    m = re.search(r'"internal_monologue"\s*:\s*"((?:[^"\\]|\\.)*)"', text)
    data["internal_monologue"] = m.group(1) if m else ""
    
    m = re.search(r'"response"\s*:\s*"((?:[^"\\]|\\.)*)"', text)
    if m:
        data["response"] = m.group(1)
    else:
        m = re.search(r'"response"\s*:\s*([^,}]+?)(?=,\s*"|}|$)', text)
        data["response"] = m.group(1).strip().strip('"').strip() if m else ""
    
    m = re.search(r'"emotion"\s*:\s*"([^"]*)"', text)
    data["emotion"] = m.group(1) if m else "?"
    
    for flag in ['red_eyes', 'is_angry', 'is_following', 'is_suspicious',
                  'is_threatening', 'open_main_door', 'DontLetPlayerLeave']:
        m = re.search(rf'"{flag}"\s*:\s*(true|false)', text)
        data[flag] = (m.group(1) == 'true') if m else False
    
    return data if data.get("response") else None

# ── Build prompt ──────────────────────────────────────────────────────────────
def build_aiko_prompt(history, user_msg, location="middle of the apartment"):
    parts = []
    for past_user, past_json in history:
        parts.append(f"[user] Location: {location} | Player user: {past_user}")
        parts.append(f"AI: {past_json}")
    
    parts.append(f"[user] Location: {location} | Player user: {user_msg}")
    return "\n".join(parts) + f"\nAI: {JSON_PREFIX}"

# ── Generation ────────────────────────────────────────────────────────────────
def _run_generation(model_instance, prompt: str, streamer: TextIteratorStreamer) -> None:
    try:
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
        with torch.inference_mode():
            model_instance.generate(
                input_ids      = inputs.input_ids,
                attention_mask = inputs.attention_mask,
                streamer       = streamer,
                **GEN_CONFIG,
            )
    except Exception as e:
        logger.error(f"Generation error: {e}")
        streamer.text_queue.put(streamer.stop_signal)
        raise

# ── Format display ────────────────────────────────────────────────────────────
def format_aiko_display(data):
    if not data:
        return "..."
    
    monologue = data.get("internal_monologue", "")
    response  = data.get("response", "...")
    emotion   = data.get("emotion", "?")
    
    parts = []
    
    if monologue:
        parts.append(f"*[{emotion}] {monologue}*")
        parts.append("")
    
    parts.append(response)
    
    flags = []
    flag_labels = {
        'red_eyes': '🔴 red_eyes',
        'is_angry': '😡 angry',
        'is_following': '👣 following',
        'is_suspicious': '🤨 suspicious',
        'is_threatening': '⚠️ threatening',
        'open_main_door': '🚪 door_open',
        'DontLetPlayerLeave': '🚫 blocking_exit',
    }
    for flag, label in flag_labels.items():
        if data.get(flag):
            flags.append(label)
    
    if flags:
        parts.append("")
        parts.append(f"`{' • '.join(flags)}`")
    
    return "\n".join(parts)

# ── Chat history storage ──────────────────────────────────────────────────────
aiko_histories = {}

def respond(message: str, history: list, api_key: str, location: str):
    if not validate_key(api_key):
        yield "unauthorized"
        return

    if not queue_counter.acquire(blocking=False):
        yield "server busy — try again"
        return

    try:
        history_key = str(id(history))
        aiko_history = aiko_histories.get(history_key, [])
        
        expected_len = len(history)
        if len(aiko_history) > expected_len:
            aiko_history = aiko_history[:expected_len]
        
        if len(aiko_history) > 6:
            aiko_history = aiko_history[-6:]
        
        prompt = build_aiko_prompt(aiko_history, message, location)

        idx, model_instance = acquire_free_model(timeout=60.0)
        if model_instance is None:
            yield "all models busy"
            return

        accumulated = ""
        try:
            streamer = TextIteratorStreamer(
                tokenizer,
                skip_prompt         = True,
                skip_special_tokens = True,
                timeout             = 120.0,
            )

            future = executor.submit(_run_generation, model_instance, prompt, streamer)

            for token in streamer:
                accumulated += token
                full_text = JSON_PREFIX + accumulated
                data = extract_aiko_data(full_text)
                if data and data.get("response"):
                    yield format_aiko_display(data)

            future.result(timeout=10)
            
            full_text = JSON_PREFIX + accumulated
            data = extract_aiko_data(full_text)
            
            if data and data.get("response"):
                try:
                    end_idx = full_text.rindex('}') + 1
                    clean_json = full_text[:end_idx]
                except:
                    clean_json = json.dumps(data)
                
                aiko_history.append((message, clean_json))
                aiko_histories[history_key] = aiko_history
                
                yield format_aiko_display(data)
            else:
                yield "..."

        except Exception as e:
            logger.error(f"respond() error: {e}")
            if not accumulated:
                yield "generation failed"
        finally:
            model_locks[idx].release()

    finally:
        queue_counter.release()


# ── Warmup ────────────────────────────────────────────────────────────────────
def _warmup():
    logger.info("Warmup starting...")
    try:
        prompt = "[user] Location: middle of the apartment | Player user: hi\nAI: " + JSON_PREFIX
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=256)
        warmup_cfg = {**GEN_CONFIG, "max_new_tokens": 20}

        for idx, m in enumerate(models):
            t0 = time.time()
            try:
                with torch.inference_mode():
                    m.generate(
                        input_ids      = inputs.input_ids,
                        attention_mask = inputs.attention_mask,
                        **warmup_cfg,
                    )
                logger.info(f"Model #{idx} warmup: {time.time()-t0:.1f}s ✓")
            except Exception as e:
                logger.warning(f"Model #{idx} warmup failed: {e}")
        logger.info("All warmups complete ✓")
    except Exception as e:
        logger.warning(f"Warmup failed: {e}")

Thread(target=_warmup, daemon=True).start()


# ── CSS — Simple Pink AI Theme ────────────────────────────────────────────────
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

* { font-family: 'Inter', sans-serif !important; }

html, body, .gradio-container, .main, .wrap, gradio-app {
    background-color: #fafafa !important;
    color: #1a1a1a !important;
}

.gradio-container {
    max-width: 800px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}

/* Headers */
h1 {
    color: #ec4899 !important;
    font-weight: 700 !important;
    font-size: 28px !important;
    text-align: center !important;
    margin: 0 !important;
}

h3 {
    color: #6b7280 !important;
    font-weight: 400 !important;
    font-size: 14px !important;
    text-align: center !important;
    margin-top: 4px !important;
}

/* Chatbot */
.chatbot, [data-testid="chatbot"], .message-wrap, .messages {
    background: #ffffff !important;
    border: 1px solid #f3e8ff !important;
    border-radius: 12px !important;
    box-shadow: 0 1px 3px rgba(236, 72, 153, 0.05) !important;
}

/* User messages — pink */
.message.user .message-bubble-border, 
div[data-testid="user"] .message-bubble-border,
.message.user {
    background: #ec4899 !important;
    color: #ffffff !important;
    border: none !important;
    border-radius: 16px 16px 4px 16px !important;
    font-size: 14px !important;
    line-height: 1.5 !important;
    padding: 10px 14px !important;
}

/* Bot messages — light pink */
.message.bot .message-bubble-border,
div[data-testid="bot"] .message-bubble-border,
.message.bot {
    background: #fdf2f8 !important;
    color: #1a1a1a !important;
    border: 1px solid #fce7f3 !important;
    border-radius: 16px 16px 16px 4px !important;
    font-size: 14px !important;
    line-height: 1.6 !important;
    padding: 12px 16px !important;
}

.message.user *, .message.user p { color: #ffffff !important; }
.message.bot *, .message.bot p { color: #1a1a1a !important; }

/* Italic monologue */
.message.bot em, .message.bot i {
    color: #ec4899 !important;
    font-style: italic;
    display: block;
    margin-bottom: 6px;
    font-size: 12px;
    opacity: 0.85;
}

/* Code (flags) */
.message.bot code {
    background: #fce7f3 !important;
    color: #be185d !important;
    padding: 3px 8px !important;
    border-radius: 6px !important;
    font-size: 11px !important;
    font-family: 'Inter', sans-serif !important;
    display: inline-block !important;
    margin-top: 6px !important;
}

/* Avatars */
.avatar-container, .avatar-container img, img.avatar-image {
    width: 36px !important;
    height: 36px !important;
    border-radius: 50% !important;
}

/* Input */
textarea, input[type="text"] {
    background: #ffffff !important;
    color: #1a1a1a !important;
    border: 1px solid #fce7f3 !important;
    border-radius: 10px !important;
    font-size: 14px !important;
    padding: 10px 14px !important;
    transition: border-color 0.2s ease !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: #ec4899 !important;
    box-shadow: 0 0 0 3px rgba(236, 72, 153, 0.1) !important;
    outline: none !important;
}

textarea::placeholder, input::placeholder {
    color: #9ca3af !important;
}

/* Buttons */
button {
    background: #ec4899 !important;
    color: #ffffff !important;
    border: none !important;
    border-radius: 8px !important;
    font-size: 13px !important;
    font-weight: 500 !important;
    padding: 8px 20px !important;
    transition: background-color 0.2s ease !important;
    cursor: pointer !important;
}

button:hover {
    background: #db2777 !important;
}

/* Lock screen */
#lock-screen {
    position: fixed !important;
    top: 0 !important; left: 0 !important;
    width: 100vw !important; height: 100vh !important;
    background: #fafafa !important;
    z-index: 9999 !important;
    display: flex !important;
    flex-direction: column !important;
    align-items: center !important;
    justify-content: center !important;
    gap: 16px !important;
}

#lock-screen h1 {
    font-size: 36px !important;
    margin-bottom: 4px !important;
}

#lock-screen input {
    width: 280px !important;
    max-width: 80vw !important;
}

#lock-screen button {
    width: 160px !important;
}

/* Misc */
footer { display: none !important; }

.block, .form, .gap {
    background: transparent !important;
    border: none !important;
}

label, .label {
    color: #6b7280 !important;
    font-size: 13px !important;
    font-weight: 500 !important;
}

/* Scrollbar */
::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: #fafafa; }
::-webkit-scrollbar-thumb { background: #ec4899; border-radius: 3px; }

/* Dropdown */
select, .gr-dropdown {
    background: #ffffff !important;
    color: #1a1a1a !important;
    border: 1px solid #fce7f3 !important;
    border-radius: 8px !important;
    padding: 8px 12px !important;
    font-size: 13px !important;
}
"""

LOCATIONS = [
    "middle of the apartment",
    "living room",
    "kitchen",
    "bedroom",
    "bathroom",
    "hallway",
    "front door",
]

# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:

    # Lock screen
    with gr.Column(elem_id="lock-screen", visible=True) as lock_screen:
        gr.Markdown("# Aiko")
        gr.Markdown("### enter access key")
        key_input  = gr.Textbox(
            placeholder="access key",
            type="password",
            show_label=False,
            max_lines=1,
        )
        unlock_btn = gr.Button("ENTER")
        lock_msg   = gr.Markdown("")

    # Chat screen
    with gr.Column(visible=False) as chat_screen:
        gr.Markdown("# Aiko")
        
        stored_key = gr.Textbox(value="", visible=False)
        
        location_dropdown = gr.Dropdown(
            choices=LOCATIONS,
            value="middle of the apartment",
            label="Location",
            interactive=True,
        )

        chatbot = gr.Chatbot(
            type              = "messages",
            height            = 500,
            show_label        = False,
            bubble_full_width = False,
            show_copy_button  = False,
        )

        gr.ChatInterface(
            fn                          = respond,
            chatbot                     = chatbot,
            type                        = "messages",
            additional_inputs           = [stored_key, location_dropdown],
            additional_inputs_accordion = gr.Accordion(visible=False),
        )

    def try_unlock(key):
        if validate_key(key):
            return gr.update(visible=False), gr.update(visible=True), key, ""
        return gr.update(visible=True), gr.update(visible=False), "", "invalid key"

    for trigger in (unlock_btn.click, key_input.submit):
        trigger(fn=try_unlock, inputs=[key_input],
                outputs=[lock_screen, chat_screen, stored_key, lock_msg])

if __name__ == "__main__":
    demo.launch(
        max_threads = MAX_CONCURRENT * 4,
        show_error  = True,
    )