app.py

import os
import spaces
import torch
import tempfile
import asyncio
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import edge_tts
import pdfplumber

MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
API_KEY  = os.environ.get("BRAIN_API_KEY", "")
TTS_VOZ  = "es-ES-AlvaroNeural"

SYSTEM_DEFAULT = (
    "Eres F\u00e9nix Brain, experto en Python, Gradio, FastAPI y HuggingFace Spaces. "
    "Responde siempre en espa\u00f1ol con c\u00f3digo limpio y funcional."
)

print("Cargando tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

print("Configurando cuantizaci\u00f3n (4-bits)...")
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

print("Cargando modelo...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map="auto",
    low_cpu_mem_usage=True,
    trust_remote_code=True,
)
print("Modelo 32B listo \u2705")


async def _tts_async(texto: str, voz: str) -> str:
    tmp = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
    await edge_tts.Communicate(texto[:3000], voz).save(tmp.name)
    return tmp.name

def texto_a_voz(texto: str, voz: str) -> str:
    try:
        return asyncio.new_event_loop().run_until_complete(_tts_async(texto, voz))
    except Exception as e:
        print(f"TTS error: {e}")
        return None


def _leer_archivo(path: str) -> str:
    if path.lower().endswith(".pdf"):
        try:
            with pdfplumber.open(path) as pdf:
                return "\n".join(p.extract_text() or "" for p in pdf.pages[:10])
        except Exception as e:
            return f"[Error leyendo PDF: {e}]"
    else:
        try:
            with open(path, "r", encoding="utf-8", errors="ignore") as f:
                return f.read()[:8000]
        except Exception as e:
            return f"[Error leyendo archivo: {e}]"


@spaces.GPU
def generar(prompt: str, system: str, key: str, max_tokens: int = 1024) -> str:
    """Endpoint API para La Forja."""
    if API_KEY and key != API_KEY:
        return "\u274c Unauthorized"
    if not prompt.strip():
        return "\u274c Prompt vac\u00edo"
    messages = [
        {"role": "system", "content": system or SYSTEM_DEFAULT},
        {"role": "user",   "content": prompt},
    ]
    text   = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs, max_new_tokens=max_tokens,
            temperature=0.2, do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    gen = out[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(gen, skip_special_tokens=True)


@spaces.GPU
def chat_fn(mensaje: str, archivo, historial: list, system_prompt: str, voz: str, activar_voz: bool):
    if not mensaje.strip() and archivo is None:
        return historial, "", None

    texto_archivo = ""
    if archivo is not None:
        contenido = _leer_archivo(archivo.name)
        nombre = os.path.basename(archivo.name)
        texto_archivo = f"\n\n[Archivo: {nombre}]\n{contenido}"

    prompt_completo = mensaje + texto_archivo

    messages = [{"role": "system", "content": system_prompt or SYSTEM_DEFAULT}]
    for h in historial:
        if isinstance(h, dict):
            messages.append(h)
        else:
            messages.append({"role": "user",      "content": h[0]})
            messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": prompt_completo})

    text   = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs, max_new_tokens=1024, temperature=0.2,
            do_sample=True, pad_token_id=tokenizer.eos_token_id,
        )
    gen = out[0][inputs["input_ids"].shape[1]:]
    respuesta = tokenizer.decode(gen, skip_special_tokens=True)

    historial = historial + [
        {"role": "user",      "content": mensaje + (" \U0001f4c4" if archivo else "")},
        {"role": "assistant", "content": respuesta},
    ]

    audio_path = texto_a_voz(respuesta, voz) if activar_voz else None
    return historial, "", audio_path


VOCES = {
    "\U0001f1ea\U0001f1f8 \u00c1lvaro (ES)": "es-ES-AlvaroNeural",
    "\U0001f1ea\U0001f1f8 Elvira (ES)":      "es-ES-ElviraNeural",
    "\U0001f1f2\U0001f1fd Jorge (MX)":       "es-MX-JorgeNeural",
    "\U0001f1f2\U0001f1fd Dalia (MX)":       "es-MX-DaliaNeural",
    "\U0001f1e6\U0001f1f7 Tom\u00e1s (AR)":  "es-AR-TomasNeural",
}

with gr.Blocks(title="\U0001f9e0 F\u00e9nix Brain", theme=gr.themes.Base()) as demo:
    gr.Markdown("# \U0001f9e0 F\u00e9nix Brain\nQwen2.5-Coder-32B-Instruct (4-bit) \u00b7 Documentos \u00b7 Voz")

    with gr.Tab("\U0001f4ac Chat"):
        with gr.Row():
            with gr.Column(scale=3):
                system_box = gr.Textbox(label="System prompt", value=SYSTEM_DEFAULT, lines=2)
                chatbot    = gr.Chatbot(height=450, type="messages", label="F\u00e9nix Brain")
                msg_box    = gr.Textbox(placeholder="Escribe tu mensaje\u2026", label="Mensaje", lines=2)
                with gr.Row():
                    send_btn  = gr.Button("\U0001f4e8 Enviar", variant="primary")
                    clear_btn = gr.Button("\U0001f5d1\ufe0f Limpiar")
            with gr.Column(scale=1):
                archivo_in  = gr.File(label="\U0001f4c4 Archivo", file_types=[".pdf", ".txt", ".py", ".md", ".json"])
                activar_voz = gr.Checkbox(label="\U0001f50a Voz activada", value=True)
                voz_sel     = gr.Dropdown(list(VOCES.keys()), value="\U0001f1ea\U0001f1f8 \u00c1lvaro (ES)", label="Voz")
                audio_out   = gr.Audio(label="\U0001f50a Respuesta", autoplay=True)

    def _chat(mensaje, archivo, historial, system, voz_nombre, activar):
        voz = VOCES.get(voz_nombre, TTS_VOZ)
        return chat_fn(mensaje, archivo, historial, system, voz, activar)

    send_btn.click(_chat, [msg_box, archivo_in, chatbot, system_box, voz_sel, activar_voz], [chatbot, msg_box, audio_out])
    msg_box.submit(_chat, [msg_box, archivo_in, chatbot, system_box, voz_sel, activar_voz], [chatbot, msg_box, audio_out])
    clear_btn.click(lambda: ([], "", None), outputs=[chatbot, msg_box, audio_out])

    # ── Endpoint API oculto para La Forja ────────────────────────────────────
    with gr.Tab("\U0001f50c API (La Forja)"):
        gr.Markdown("### Endpoint `/generar` para La Forja")
        _p   = gr.Textbox(label="prompt",     visible=False)
        _s   = gr.Textbox(label="system",     visible=False)
        _k   = gr.Textbox(label="key",        visible=False)
        _mt  = gr.Number( label="max_tokens", visible=False, value=1024)
        _out = gr.Textbox(label="response",   visible=False)
        _btn = gr.Button("generar",           visible=False)
        _btn.click(generar, inputs=[_p, _s, _k, _mt], outputs=[_out], api_name="generar")
        gr.Markdown("""
```python
from gradio_client import Client
client = Client("Cristobal299/Chat")
result = client.predict(
    prompt="Tu pregunta",
    system="",
    key="TU_BRAIN_API_KEY",
    max_tokens=1024,
    api_name="/generar"
)
print(result)
```
""")

demo.launch(server_name="0.0.0.0", server_port=7860)