app.py

"""
╔══════════════════════════════════════════════════════════════╗
║          CYGNIS ALPHA 2 — SOVEREIGN AI INTERFACE             ║
║          Production-Grade Gradio Space · v2.0.0              ║
╚══════════════════════════════════════════════════════════════╝
Auteur  : Simonc-44
Modèle  : Cygnis-Alpha-2-8B (LoRA adapter sur Llama 3.1)
"""

import os
import sys
import time
import logging
import gc
from threading import Thread

import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from peft import PeftModel
from huggingface_hub import login

# ─────────────────────────────────────────────────────────────
# LOGGING & AUTH
# ─────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("cygnis")

hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)

# ─────────────────────────────────────────────────────────────
# CONFIGURATION
# ─────────────────────────────────────────────────────────────
BASE_MODEL_ID = "unsloth/meta-llama-3.1-8b-bnb-4bit"
ADAPTER_ID    = "Simonc-44/Cygnis-Alpha-2-8B-v0.1"

SYSTEM_PROMPT = """### IDENTITY
Vous êtes Cygnis-Alpha-2-8B, un LLM souverain conçu par Simonc-44.
Structurez vos réponses avec [RÉFLEXION], [DÉMONSTRATION] et [CONCLUSION]."""

# ─────────────────────────────────────────────────────────────
# CHARGEMENT (OPTIMISÉ)
# ─────────────────────────────────────────────────────────────
def load_model():
    logger.info("🚀 Chargement du tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
    tokenizer.pad_token = tokenizer.eos_token

    logger.info("📦 Chargement du modèle de base...")
    # device_map="auto" place automatiquement sur GPU si dispo, sinon CPU
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        device_map="auto",
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )

    logger.info("💉 Application de l'adapter LoRA...")
    model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
    model.eval()
    
    logger.info("✅ Cygnis est prêt !")
    return tokenizer, model

try:
    tokenizer, model = load_model()
    MODEL_LOADED = True
except Exception as e:
    logger.error(f"❌ Échec du chargement : {e}")
    MODEL_LOADED = False

# ─────────────────────────────────────────────────────────────
# LOGIQUE D'INFÉRENCE
# ─────────────────────────────────────────────────────────────
def chat_cygnis(message, history, system_prompt, max_new_tokens, temperature):
    if not MODEL_LOADED:
        yield "Erreur : Le modèle n'a pas pu être chargé."
        return

    # Format Gradio 5 (list of dicts) vers Chat Template
    messages = [{"role": "system", "content": system_prompt}]
    for turn in history:
        messages.append(turn)
    messages.append({"role": "user", "content": message})

    # Préparation des inputs
    input_ids = tokenizer.apply_chat_template(
        messages, 
        add_generation_prompt=True, 
        return_tensors="pt"
    ).to(model.device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
        do_sample=(float(temperature) > 0),
        pad_token_id=tokenizer.eos_token_id,
    )

    # Lancement du thread de génération
    thread = Thread(target=model.generate, kwargs=generate_kwargs)
    thread.start()

    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

# ─────────────────────────────────────────────────────────────
# INTERFACE GRADIO (SANS CSS)
# ─────────────────────────────────────────────────────────────
with gr.Blocks(title="Cygnis Alpha 2") as demo:
    gr.Markdown("# CYGNIS ALPHA 2")
    gr.Markdown("Sovereign Reasoning Engine v0.3")
    
    chatbot = gr.Chatbot(height=500, type="messages")
    msg_input = gr.Textbox(placeholder="Posez votre question ici...", label="Entrée")
    
    with gr.Accordion("Paramètres de génération", open=False):
        sys_prompt = gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt", lines=5)
        temp = gr.Slider(0.1, 1.5, value=0.3, label="Température")
        tokens = gr.Slider(64, 1024, value=512, label="Max New Tokens")

    def respond(message, chat_history, sys_p, t, tok):
        # Ajouter le message utilisateur à l'historique
        chat_history.append({"role": "user", "content": message})
        # Envoyer l'historique mis à jour (vide pour l'assistant pour l'instant)
        yield "", chat_history
        
        # Lancer l'inférence
        response = ""
        chat_history.append({"role": "assistant", "content": ""})
        
        for partial in chat_cygnis(message, chat_history[:-2], sys_p, tok, t):
            response = partial
            chat_history[-1]["content"] = response
            yield "", chat_history

    # Événements
    msg_input.submit(respond, [msg_input, chatbot, sys_prompt, temp, tokens], [msg_input, chatbot])

# ─────────────────────────────────────────────────────────────
# LANCEMENT
# ─────────────────────────────────────────────────────────────
if __name__ == "__main__":
    # Le paramètre queue() est vital pour le streaming
    demo.queue().launch(show_error=True)